closes #16

Signed-off-by: Valerian Saliou <valerian@valeriansaliou.name>

src/executor/count.rs

@@ -4,32 +4,63 @@
 // Copyright: 2019, Valerian Saliou <valerian@valeriansaliou.name>
 // License: Mozilla Public License v2.0 (MPL v2.0)
 
+use crate::store::fst::{StoreFSTActionBuilder, StoreFSTMisc};
+use crate::store::fst::{StoreFSTPool, GRAPH_ACCESS_LOCK};
 use crate::store::item::StoreItem;
+use crate::store::kv::StoreKVActionBuilder;
 use crate::store::kv::{StoreKVPool, STORE_ACCESS_LOCK};
 
 pub struct ExecutorCount;
 
 impl ExecutorCount {
     pub fn execute<'a>(store: StoreItem<'a>) -> Result<u32, ()> {
         match store {
-            StoreItem(collection, _bucket_value, _object_value) => {
+            // Count terms in (collection, bucket, object) from KV
+            StoreItem(collection, Some(bucket), Some(object)) => {
                 // Important: acquire database access read lock, and reference it in context. This \
                 //   prevents the database from being erased while using it in this block.
                 let _kv_access = STORE_ACCESS_LOCK.read().unwrap();
 
-                if let Ok(_kv_store) = StoreKVPool::acquire(collection) {
-                    // let kv_action = StoreKVActionBuilder::read(bucket, kv_store);
+                if let Ok(kv_store) = StoreKVPool::acquire(collection) {
+                    let kv_action = StoreKVActionBuilder::read(bucket, kv_store);
 
-                    // TODO: if object, count terms in object (from kv directly)
-                    // TODO: if bucket, count terms (from fst directly)
-                    // TODO: if collection, count buckets (from fs directly)
+                    // Try to resolve existing OID to IID
+                    let oid = object.as_str().to_owned();
 
-                    // TODO
-                    return Ok(0);
+                    kv_action.get_oid_to_iid(&oid).unwrap_or(None)
+                        .map(|iid| {
+                            // List terms for IID
+                            if let Some(terms) = kv_action.get_iid_to_terms(iid).unwrap_or(None) {
+                                terms.len() as u32
+                            } else {
+                                0
+                            }
+                        })
+                        .ok_or(())
+                        .or(Ok(0))
+                } else {
+                    Err(())
                 }
             }
-        }
+            // Count terms in (collection, bucket) from FST
+            StoreItem(collection, Some(bucket), None) => {
+                // Important: acquire graph access read lock, and reference it in context. This \
+                //   prevents the graph from being erased while using it in this block.
+                let _fst_access = GRAPH_ACCESS_LOCK.read().unwrap();
+
+                if let Ok(fst_store) = StoreFSTPool::acquire(collection, bucket) {
+                    let fst_action = StoreFSTActionBuilder::read(fst_store);
 
-        Err(())
+                    Ok(fst_action.count_words() as u32)
+                } else {
+                    Err(())
+                }
+            }
+            // Count buckets in (collection) from FS
+            StoreItem(collection, None, None) => {
+                StoreFSTMisc::count_collection_buckets(collection).map(|count| count as u32)
+            }
+            _ => Err(()),
+        }
     }
 }

src/store/fst.rs

@@ -48,6 +48,8 @@ pub struct StoreFSTAction {
     store: StoreFSTBox,
 }
 
+pub struct StoreFSTMisc;
+
 #[derive(Copy, Clone)]
 enum StoreFSTPathMode {
     Permanent,
@@ -485,6 +487,10 @@ impl StoreFSTBuilder {
 }
 
 impl StoreFST {
+    pub fn cardinality(&self) -> usize {
+        self.graph.len()
+    }
+
     pub fn contains(&self, word: &str) -> bool {
         let word_bytes = word.as_bytes();
 
@@ -650,7 +656,10 @@ impl StoreFSTActionBuilder {
                     }
                 }
             } else {
-                warn!("failed reading directory: {:?}", collection_path);
+                warn!(
+                    "failed reading directory for erasure: {:?}",
+                    collection_path
+                );
             }
 
             // Force a FST graph close (on all contained buckets)
@@ -830,6 +839,10 @@ impl StoreFSTAction {
         }
     }
 
+    pub fn count_words(&self) -> usize {
+        self.store.cardinality()
+    }
+
     fn find_words_stream<A: Automaton>(
         mut stream: FSTStream<A>,
         found_words: &mut Vec<String>,
@@ -851,3 +864,43 @@ impl StoreFSTAction {
         }
     }
 }
+
+impl StoreFSTMisc {
+    pub fn count_collection_buckets<'a, T: Into<&'a str>>(collection: T) -> Result<usize, ()> {
+        let collection_str = collection.into();
+
+        let mut count = 0;
+
+        let path_mode = StoreFSTPathMode::Permanent;
+        let collection_path = StoreFSTBuilder::path(path_mode, collection_str, None);
+
+        if collection_path.exists() == true {
+            // Scan collection directory for contained buckets (count them)
+            if let Ok(entries) = fs::read_dir(&collection_path) {
+                let fst_extension = path_mode.extension();
+                let fst_extension_len = fst_extension.len();
+
+                for entry in entries {
+                    if let Ok(entry) = entry {
+                        if let Some(entry_name) = entry.file_name().to_str() {
+                            let entry_name_len = entry_name.len();
+
+                            // FST file found? This is a bucket.
+                            if entry_name_len > fst_extension_len
+                                && entry_name.ends_with(fst_extension) == true
+                            {
+                                count += 1;
+                            }
+                        }
+                    }
+                }
+            } else {
+                warn!("failed reading directory for count: {:?}", collection_path);
+
+                return Err(());
+            }
+        }
+
+        Ok(count)
+    }
+}