pocketpy/src/objects/dict.cpp

#include "pocketpy/objects/dict.hpp"

namespace pkpy{

    Dict::Dict(): _capacity(__Capacity),
            _mask(__Capacity-1),
            _size(0), _critical_size(__Capacity*__LoadFactor+0.5f), _head_idx(-1), _tail_idx(-1){
        __alloc_items();
    }

    void Dict::__alloc_items(){
        _items = (Item*)std::malloc(_capacity * sizeof(Item));
        for(int i=0; i<_capacity; i++){
            _items[i].first = nullptr;
            _items[i].second = nullptr;
            _items[i].prev = -1;
            _items[i].next = -1;
        }
    }

    Dict::Dict(Dict&& other){
        _capacity = other._capacity;
        _mask = other._mask;
        _size = other._size;
        _critical_size = other._critical_size;
        _head_idx = other._head_idx;
        _tail_idx = other._tail_idx;
        _items = other._items;
        other._items = nullptr;
    }

    Dict::Dict(const Dict& other){
        _capacity = other._capacity;
        _mask = other._mask;
        _size = other._size;
        _critical_size = other._critical_size;
        _head_idx = other._head_idx;
        _tail_idx = other._tail_idx;
        // copy items
        _items = (Item*)std::malloc(_capacity * sizeof(Item));
        std::memcpy(_items, other._items, _capacity * sizeof(Item));
    }

    void Dict::set(VM* vm, PyVar key, PyVar val){
        // do possible rehash
        if(_size+1 > _critical_size) _rehash(vm);
        bool ok; int i;
        _probe_1(vm, key, ok, i);
        if(!ok) {
            _size++;
            _items[i].first = key;

            // append to tail
            if(_size == 0+1){
                _head_idx = i;
                _tail_idx = i;
            }else{
                _items[i].prev = _tail_idx;
                _items[_tail_idx].next = i;
                _tail_idx = i;
            }
        }
        _items[i].second = val;
    }

    void Dict::_rehash(VM* vm){
        Item* old_items = _items;
        int old_head_idx = _head_idx;

        _capacity *= 4;
        _mask = _capacity - 1;
        _size = 0;
        _critical_size = _capacity*__LoadFactor+0.5f;
        _head_idx = -1;
        _tail_idx = -1;
        
        __alloc_items();

        // copy old items to new dict
        int i = old_head_idx;
        while(i != -1){
            set(vm, old_items[i].first, old_items[i].second);
            i = old_items[i].next;
        }

        std::free(old_items);
    }


    PyVar Dict::try_get(VM* vm, PyVar key) const{
        bool ok; int i;
        _probe_0(vm, key, ok, i);
        if(!ok) return nullptr;
        return _items[i].second;
    }

    bool Dict::contains(VM* vm, PyVar key) const{
        bool ok; int i;
        _probe_0(vm, key, ok, i);
        return ok;
    }

    bool Dict::del(VM* vm, PyVar key){
        bool ok; int i;
        _probe_0(vm, key, ok, i);
        if(!ok) return false;
        _items[i].first = nullptr;
        // _items[i].second = PY_DELETED_SLOT;  // do not change .second if it is not NULL, it means the slot is occupied by a deleted item
        _size--;

        if(_size == 0){
            _head_idx = -1;
            _tail_idx = -1;
        }else{
            if(_head_idx == i){
                _head_idx = _items[i].next;
                _items[_head_idx].prev = -1;
            }else if(_tail_idx == i){
                _tail_idx = _items[i].prev;
                _items[_tail_idx].next = -1;
            }else{
                _items[_items[i].prev].next = _items[i].next;
                _items[_items[i].next].prev = _items[i].prev;
            }
        }
        _items[i].prev = -1;
        _items[i].next = -1;
        return true;
    }

    void Dict::update(VM* vm, const Dict& other){
        other.apply([&](PyVar k, PyVar v){ set(vm, k, v); });
    }

    Tuple Dict::keys() const{
        Tuple t(_size);
        int i = _head_idx;
        int j = 0;
        while(i != -1){
            t[j++] = _items[i].first;
            i = _items[i].next;
        }
        assert(j == _size);
        return t;
    }

    Tuple Dict::values() const{
        Tuple t(_size);
        int i = _head_idx;
        int j = 0;
        while(i != -1){
            t[j++] = _items[i].second;
            i = _items[i].next;
        }
        assert(j == _size);
        return t;
    }

    void Dict::clear(){
        _size = 0;
        _head_idx = -1;
        _tail_idx = -1;
        for(int i=0; i<_capacity; i++){
            _items[i].first = nullptr;
            _items[i].second = nullptr;
            _items[i].prev = -1;
            _items[i].next = -1;
        }
    }

    Dict::~Dict(){
        if(_items) std::free(_items);
    }
}   // namespace pkpy