1 files changed, 59 insertions, 21 deletions

diff --git a/database/src/kv/mod.rs b/database/src/kv/mod.rs
index 1a058a7..2d2f3cc 100644
--- a/database/src/kv/mod.rs
+++ b/database/src/kv/mod.rs
@@ -17,8 +17,8 @@ use crate::{
     Database, Query, RowIter, RowNum, Transaction,
     kv::{
         helpers::{read_counter, write_counter},
-        index::update_index,
-        index_key::IndexKey,
+        index::{read_count_index, update_index},
+        index_key::{IndexKey, SortKey},
         prefix_iterator::PrefixIterator,
     },
 };
@@ -54,9 +54,8 @@ impl Transaction for &mut dyn jellykv::Transaction {
         )?;
 
         let ob = entry.as_object();
-
         for (is, ik) in list_indices(*self)? {
-            update_index(*self, is, ik, row, ob, false)?;
+            update_index(*self, is, &ik, row, ob, false)?;
         }
 
         Ok(row)
@@ -64,9 +63,9 @@ impl Transaction for &mut dyn jellykv::Transaction {
     fn remove(&mut self, row: RowNum) -> Result<()> {
         let entry = self.get(row)?.ok_or(anyhow!("row did not exist"))?;
         let ob = entry.as_object();
-        // for index in self.indices.values() {
-        //     index.remove(db, row, ob)?;
-        // }
+        for (is, ik) in list_indices(*self)? {
+            update_index(*self, is, &ik, row, ob, true)?;
+        }
         jellykv::Transaction::del(*self, &row_key(row))
     }
     fn update(&mut self, row: RowNum, entry: ObjectBuffer) -> Result<()> {
@@ -80,12 +79,10 @@ impl Transaction for &mut dyn jellykv::Transaction {
             bytemuck::cast_slice(entry.0.as_slice()),
         )?;
 
-        // for index in self.indices.values() {
-        //     if !index.compare(before, after) {
-        //         index.remove(txn, row, before)?;
-        //         index.add(txn, row, after)?;
-        //     }
-        // }
+        for (is, ik) in list_indices(*self)? {
+            update_index(*self, is, &ik, row, before, true)?;
+            update_index(*self, is, &ik, row, after, false)?;
+        }
 
         Ok(())
     }
@@ -94,7 +91,7 @@ impl Transaction for &mut dyn jellykv::Transaction {
         Ok(jellykv::Transaction::get(*self, &row_key(row))?.map(ObjectBuffer::from))
     }
 
-    fn query(&self, query: Query) -> Result<RowIter> {
+    fn query(&mut self, query: Query) -> Result<RowIter> {
         // query
         //     .filter
         //     .get_bins()
@@ -106,20 +103,58 @@ impl Transaction for &mut dyn jellykv::Transaction {
         //     .map(|i| i.query(txn, &query.sort))
         todo!()
     }
-    fn query_single(&self, query: Query) -> Result<Option<RowNum>> {
+    fn query_single(&mut self, query: Query) -> Result<Option<RowNum>> {
         self.query(query)?.next().transpose()
     }
-    fn count(&self, query: Query) -> Result<u64> {
+    fn count(&mut self, query: Query) -> Result<u64> {
         let mut total = 0;
-        // for b in query.filter.get_bins() {
-        //     if let Some(&c) = self.counters.get(&b) {
-        //         total += read_counter(txn, c)?;
-        //     }
-        // }
+        for (binning, prefix) in query.filter.get_bins() {
+            let ik = IndexKey(binning, SortKey::Count);
+            let is = get_or_create_index(*self, &ik)?;
+            total += read_count_index(*self, is, prefix)?;
+        }
         Ok(total)
     }
 }
 
+fn get_or_create_index(txn: &mut dyn jellykv::Transaction, ik: &IndexKey) -> Result<SubtreeNum> {
+    match jellykv::Transaction::get(txn, &ik.to_bytes())? {
+        Some(is) => Ok(SubtreeNum::from_be_bytes(is.try_into().unwrap())),
+        None => {
+            let is = alloc_subtree(txn)?;
+            create_index(txn, is, ik)?;
+            jellykv::Transaction::set(txn, &ik.to_bytes(), &is.to_be_bytes())?;
+            Ok(is)
+        }
+    }
+}
+
+fn alloc_subtree(txn: &mut dyn jellykv::Transaction) -> Result<SubtreeNum> {
+    let is = read_counter(
+        txn,
+        &T_INDEX_COUNTER.to_be_bytes(),
+        INDEX_TNUM_OFFSET as u64,
+    )?;
+    write_counter(txn, &T_INDEX_COUNTER.to_be_bytes(), is + 1)?;
+    Ok(is as SubtreeNum)
+}
+
+fn create_index(txn: &mut dyn jellykv::Transaction, is: SubtreeNum, ik: &IndexKey) -> Result<()> {
+    let rowkeys = PrefixIterator {
+        inner: txn.iter(&T_ROWS.to_be_bytes(), false)?,
+        prefix: Cow::Borrowed(&T_ROWS.to_be_bytes()),
+    }
+    .collect::<Result<Vec<_>, _>>()?; // TODO dont collect this
+
+    for rowkey in rowkeys {
+        let row = inv_row_key(&rowkey);
+        let buf = ObjectBuffer::from(jellykv::Transaction::get(txn, &rowkey)?.unwrap());
+        update_index(txn, is, ik, row, buf.as_object(), false)?;
+    }
+
+    Ok(())
+}
+
 fn list_indices(txn: &dyn jellykv::Transaction) -> Result<Vec<(SubtreeNum, IndexKey)>> {
     let indices = PrefixIterator {
         inner: jellykv::Transaction::iter(txn, &T_INDICES.to_be_bytes(), false)?,
@@ -141,6 +176,9 @@ fn row_key(row: RowNum) -> Vec<u8> {
     key.extend(row.to_be_bytes());
     key
 }
+fn inv_row_key(k: &[u8]) -> RowNum {
+    RowNum::from_be_bytes(k[4..12].try_into().unwrap())
+}
 
 #[cfg(test)]
 mod test {