mdb.mo

import Map "mo:motoko-hash-map/Map";
import Set "mo:motoko-hash-map/Set";
import Blob "mo:base/Blob";
import BTree "mo:stableheapbtreemap/BTree";
import Nat64 "mo:base/Nat64";
import Iter "mo:base/Iter";
import Text "mo:base/Text";
import E "mo:lexicographic-encoding/EncodeInt";
import Array "mo:base/Array";
import Debug "mo:base/Debug";
import Vector "mo:vector";


module {
    let { thash } = Map;
    let enc = E.encodeInt;

    public type IndexKey = Text; //[Nat8];
    public type Direction = BTree.Direction;
    let IndexCompare = Text.compare;
    type BTreeOrder = Nat;
    // public type KeyFunc<K> = (K -> Nat32, (K, K) -> Bool, () -> K);
    public type IndexFunc<V> = (V, Nat, IndexOperation) -> [(BTreeOrder, Text, IndexKey)];
    public type IndexStore = BTree.BTree<IndexKey, Nat>;
    public type Index<V> = Vector.Vector<(Text, IndexStore)>;   
  
    public type IndexOperation = {
        #insert;
        #remove;
    };

    public type MDB< V> = {
        vec: Vector.Vector<?V>;
        idx: Index<V>;
    };

    public type Config<K, V> = {
        pkey_select: (V -> K);
        pkey_encode: (K -> IndexKey);
    };

    public func Init<K, V>(conf : Config<K, V>) : MDB<V> {
        let db : MDB<V> = {
            vec = Vector.new<?V>();
            idx = Vector.new<(Text, IndexStore)>();
        };
        // Vector.add(db.idx, ("primary", BTree.init<IndexKey, Nat>(?32)));
        return db;
    };

    public type MDBActions<K, V> = {
        insert: (V) -> ();
        setByIdx: (Nat, V) -> ();
        getIdx: (K) -> ?Nat;
        get: (K) -> ?V;
        getByIdx: (Nat) -> ?V;
        removeByIdx: (Nat) -> Bool;
        remove: (K) -> Bool;
        find: (Text, Direction, IndexKey, IndexKey, Nat) -> [V];
        size: () -> Nat;
        indexConfig: (IndexFunc<V>) -> ();
        indexSize: (Text) -> Nat;
        indexMin: (Text) -> (?(IndexKey, Nat));
        indexRemove: (Text, IndexKey) -> ();
    };


    public func Use<K, V>(db: MDB<V>, conf: Config<K, V>) : MDBActions<K, V> {

        var _indexing : IndexFunc<V> = func (x, idx, op) {
            [];
        };
        let indexConfig = func(fh: IndexFunc<V>) : () {
            _indexing := fh;
        };

        let get_index = func (name: Text) : ?IndexStore {
            let vals = Vector.vals(db.idx);
            for ((idxname, idx) in vals) {
                if (idxname == name) {
                    return ?idx;
                };
            };
            null;
        };

        

        let remove_idx_entries = func (id: Nat, v:V) {
             // remove old index entries
            let oldidxs = _indexing(v, id, #remove);
            let oldidxiter = Iter.fromArray(oldidxs);
            for ((btorder, idxname, idxkey) in oldidxiter) {
                let ?idx = get_index(idxname) else Debug.trap("IE3 internal error");
                ignore BTree.delete(idx, IndexCompare, idxkey);
                };
        };

        let insert_idx_entries = func (id: Nat, v:V) {
            // insert new index entries
            let newidxs = _indexing(v, id, #insert);
            let newidxiter = Iter.fromArray(newidxs);
            for ((btorder, idxname, idxkey) in newidxiter) {
                switch(get_index(idxname)) {
                    case (?idx) {
                        ignore BTree.insert(idx, IndexCompare, idxkey # "." # enc(id), id);
                    };
                    case (null) {
                        let idx = BTree.init<IndexKey, Nat>(?btorder);
                        ignore BTree.insert(idx, IndexCompare, idxkey # "." # enc(id), id);
                        Vector.add(db.idx, (idxname, idx));
                    };
                };
            };
        };

        let setByIdx = func (idx:Nat, v:V) {
            let ?pv = Vector.get(db.vec, idx) else Debug.trap("E1 Not Found");
            remove_idx_entries(idx, pv);
            Vector.put(db.vec, idx, ?v);
            insert_idx_entries(idx, v);
        };

        let insert = func(v:V) {
            // First search if it exists (by pk)
            let id:Nat = switch(getIdx(conf.pkey_select(v))) {
                case (?cidx) {
                    let ?pv = Vector.get(db.vec, cidx) else Debug.trap("IE2 internal error");
                    remove_idx_entries(cidx, pv);
                    Vector.put(db.vec, cidx, ?v);
                    cidx;
                    };
                case (null) { // If not exists add
                    Vector.add(db.vec, ?v);
                    Vector.size(db.vec) - 1;
                }
            };
        
            insert_idx_entries(id, v);
        };

        let getIdx = func(pk:K) : ?Nat {
            let ?pk_idx = get_index("primary") else return null;
            BTree.get(pk_idx, IndexCompare, conf.pkey_encode(pk));
        };

        let get = func(pk:K) : ?V {
            let ?idx = getIdx(pk) else return null;

            getByIdx(idx);
        };

        let getByIdx = func(k:Nat) : ?V {
            Vector.get(db.vec, k);
        };

        let remove = func(pk: K) : Bool {
            let ?idx = getIdx(pk) else return false;
            removeByIdx(idx);
        };

        let removeByIdx = func(idx: Nat) : Bool {
            // remove index entries
            let ?pv = Vector.get(db.vec, idx) else return false;
            remove_idx_entries(idx, pv);
            Vector.put(db.vec, idx, null);
            true;
        };

        let find = func (indexName: Text, dir: Direction, start:IndexKey, end:IndexKey, limit: Nat) : [V] {
            let ?myidx = get_index(indexName) else Debug.trap("E3 No such index");
            let res = BTree.scanLimit<IndexKey, Nat>(myidx, IndexCompare, start, end, dir, limit);
            Array.mapEntries<(IndexKey, Nat), V>(res.results, func ((_, vecidx), idx) : (V) {
                let ?v = Vector.get(db.vec, vecidx) else Debug.trap("E10 internal error");
                v;
            });
        };

        let size = func () : Nat {
            Vector.size(db.vec);
        };

        let indexSize = func (name:Text) : Nat {
            let ?idx = get_index(name) else return 0;
            BTree.size(idx);
        };

        let indexMin = func (name:Text) : (?(IndexKey, Nat)) {
            let ?idx = get_index(name) else return null;
            BTree.min(idx);
        };

        let indexRemove = func (name:Text, key: IndexKey) : () {
            let ?idx = get_index(name) else return ();
            ignore BTree.delete(idx, IndexCompare, key);
        };

        return {indexRemove; indexMin; indexSize; insert; get; setByIdx; size; getIdx; getByIdx; removeByIdx; remove; find; indexConfig};
    }
    


}