Merge pull request #60 from LiskHQ/57_add_interface_current_state

Add interface to get current root and version - Closes #57

database.js

@@ -331,4 +331,4 @@ module.exports = {
     Database,
     InMemoryDatabase,
     Batch,
-};
+};

src/common_db.rs

@@ -170,7 +170,7 @@ impl NewDBWithContext for DB {
 impl Kind {
     pub fn key(&self, key: Vec<u8>) -> Vec<u8> {
         match self {
-            Kind::State => [consts::PREFIX_STATE, &key].concat(),
+            Kind::State => [consts::Prefix::STATE, &key].concat(),
             _ => key,
         }
     }

src/consts.rs

@@ -1,10 +1,14 @@
 use crate::types::{KeyLength, SubtreeHeight, SubtreeHeightKind};
 
-pub const PREFIX_STATE: &[u8] = &[0];
-pub const PREFIX_SMT: &[u8] = &[1];
-pub const PREFIX_DIFF: &[u8] = &[2];
-
 pub const KEY_LENGTH: KeyLength = KeyLength(38);
 pub const SUBTREE_HEIGHT: SubtreeHeight = SubtreeHeight(SubtreeHeightKind::Eight);
 
 pub static PREFIX_BRANCH_HASH: &[u8] = &[1];
+
+pub struct Prefix;
+impl Prefix {
+    pub const STATE: &'static [u8] = &[0];
+    pub const SMT: &'static [u8] = &[1];
+    pub const DIFF: &'static [u8] = &[2];
+    pub const CURRENT_STATE: &'static [u8] = &[3];
+}

src/lib.rs

@@ -45,6 +45,10 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
         "state_db_new",
         state_db::StateDB::js_new_with_box_ref::<common_db::Options, state_db::StateDB>,
     )?;
+    cx.export_function(
+        "state_db_get_current_state",
+        state_db::StateDB::js_get_current_state,
+    )?;
     cx.export_function("state_db_close", state_db::StateDB::js_close)?;
     cx.export_function("state_db_get", state_db::StateDB::js_get)?;
     cx.export_function("state_db_exists", state_db::StateDB::js_exists)?;

src/smt_db.rs

@@ -14,7 +14,7 @@ pub struct InMemorySmtDB {
 
 impl Actions for SmtDB<'_> {
     fn get(&self, key: &[u8]) -> Result<VecOption, rocksdb::Error> {
-        let result = self.db.get([consts::PREFIX_SMT, key].concat())?;
+        let result = self.db.get([consts::Prefix::SMT, key].concat())?;
         Ok(result)
     }
 

src/state_db.rs

@@ -1,8 +1,10 @@
-use neon::prelude::*;
-use neon::types::buffer::TypedArray;
 use std::cmp;
+use std::convert::TryInto;
 use std::sync::{mpsc, Arc, Mutex, MutexGuard};
 use std::thread;
+
+use neon::prelude::*;
+use neon::types::buffer::TypedArray;
 use thiserror::Error;
 
 use crate::batch;
@@ -15,7 +17,9 @@ use crate::options;
 use crate::smt;
 use crate::smt_db;
 use crate::state_writer;
-use crate::types::{ArcMutex, Height, KVPair, KeyLength, NestedVec, SharedVec};
+use crate::types::{
+    ArcMutex, BlockHeight, HashKind, HashWithKind, Height, KVPair, KeyLength, NestedVec, SharedVec,
+};
 use crate::utils;
 
 type SharedStateDB = JsBoxRef<StateDB>;
@@ -28,6 +32,12 @@ pub enum DataStoreError {
     DiffNotFound(usize),
 }
 
+#[derive(Debug, PartialEq, Eq)]
+struct CurrentState<'a> {
+    root: &'a [u8],
+    version: BlockHeight,
+}
+
 struct Commit {
     db_options: DBOptions,
     check_expected: bool,
@@ -49,6 +59,23 @@ pub struct StateDB {
     options: DBOptions,
 }
 
+impl<'a> CurrentState<'a> {
+    fn new(root: &'a [u8], version: BlockHeight) -> Self {
+        Self { root, version }
+    }
+
+    fn to_bytes(&self) -> Vec<u8> {
+        [self.root, &self.version.to_be_bytes()].concat()
+    }
+
+    fn from_bytes(bytes: &'a [u8]) -> Self {
+        let version_point = bytes.len() - 4;
+        let root = &bytes[..version_point];
+        let version = u32::from_be_bytes(bytes[version_point..].try_into().unwrap()).into();
+        Self { root, version }
+    }
+}
+
 impl Commit {
     fn new(expected: Vec<u8>, db_options: DBOptions, check_expected: bool) -> Self {
         Self {
@@ -98,7 +125,7 @@ impl StateDB {
         key_length: KeyLength,
     ) -> Result<SharedVec, DataStoreError> {
         let diff_bytes = conn
-            .get(&[consts::PREFIX_DIFF, &height.as_u32_to_be_bytes()].concat())
+            .get(&[consts::Prefix::DIFF, &height.as_u32_to_be_bytes()].concat())
             .map_err(|err| DataStoreError::Unknown(err.to_string()))?
             .ok_or_else(|| DataStoreError::DiffNotFound(height.into()))?;
 
@@ -113,13 +140,13 @@ impl StateDB {
 
         let mut write_batch = batch::PrefixWriteBatch::new();
         // Insert state batch with diff
-        write_batch.set_prefix(&consts::PREFIX_STATE);
+        write_batch.set_prefix(&consts::Prefix::STATE);
         d.revert_commit(&mut write_batch);
-        write_batch.set_prefix(&consts::PREFIX_DIFF);
+        write_batch.set_prefix(&consts::Prefix::DIFF);
         write_batch.delete(&height.as_u32_to_be_bytes());
 
         // insert SMT batch
-        write_batch.set_prefix(&consts::PREFIX_SMT);
+        write_batch.set_prefix(&consts::Prefix::SMT);
         smtdb.batch.iterate(&mut write_batch);
         // insert diff
         conn.write(write_batch.batch)
@@ -137,6 +164,12 @@ impl StateDB {
         let key_length = self.options.key_length;
         self.common.send(move |conn, channel| {
             let result = StateDB::get_revert_result(conn, height, &state_root, key_length);
+            if result.is_ok() {
+                let value = (**result.as_ref().unwrap().lock().unwrap()).clone();
+                let state_info = CurrentState::new(&value, height.sub(1).into());
+                conn.put(consts::Prefix::CURRENT_STATE, state_info.to_bytes())
+                    .expect("Update state info should not be failed");
+            }
             channel.send(move |mut ctx| {
                 let callback = cb.into_inner(&mut ctx);
                 let this = ctx.undefined();
@@ -176,19 +209,26 @@ impl StateDB {
         // Create global batch
         let mut write_batch = batch::PrefixWriteBatch::new();
         // Insert state batch with diff
-        write_batch.set_prefix(&consts::PREFIX_STATE);
+        write_batch.set_prefix(&consts::Prefix::STATE);
         let diff = writer.commit(&mut write_batch);
-        write_batch.set_prefix(&consts::PREFIX_DIFF);
+        write_batch.set_prefix(&consts::Prefix::DIFF);
         let key = info.data.db_options.key_length.as_u32_to_be_bytes();
         write_batch.put(&key, diff.encode().as_ref());
 
         // insert SMT batch
-        write_batch.set_prefix(&consts::PREFIX_SMT);
+        write_batch.set_prefix(&consts::Prefix::SMT);
         smtdb.batch.iterate(&mut write_batch);
         // insert diff
         let result = conn.write(write_batch.batch);
+        let height = info.data.db_options.key_length.into();
         match result {
-            Ok(_) => Ok(root),
+            Ok(_) => {
+                let value = (**root.as_ref().lock().unwrap()).clone();
+                let state_info = CurrentState::new(&value, height);
+                conn.put(consts::Prefix::CURRENT_STATE, state_info.to_bytes())
+                    .expect("Update state info should not be failed");
+                Ok(root)
+            },
             Err(err) => Err(smt::SMTError::Unknown(err.to_string())),
         }
     }
@@ -289,8 +329,8 @@ impl StateDB {
         self.common
             .send(move |conn, channel| {
                 let zero: u32 = 0;
-                let start = [consts::PREFIX_DIFF, zero.to_be_bytes().as_slice()].concat();
-                let end = [consts::PREFIX_DIFF, &height.sub(1).as_u32_to_be_bytes()].concat();
+                let start = [consts::Prefix::DIFF, zero.to_be_bytes().as_slice()].concat();
+                let end = [consts::Prefix::DIFF, &height.sub(1).as_u32_to_be_bytes()].concat();
                 let mut batch = rocksdb::WriteBatch::default();
 
                 let iter = conn.iterator(rocksdb::IteratorMode::From(
@@ -380,6 +420,40 @@ impl StateDB {
 
         Ok(parsed_query_keys)
     }
+
+    fn get_current_state(
+        &self,
+        cb: Root<JsFunction>,
+    ) -> Result<(), mpsc::SendError<options::DbMessage>> {
+        self.common.send(move |conn, channel| {
+            let result = conn.get(consts::Prefix::CURRENT_STATE);
+            channel.send(move |mut ctx| {
+                let callback = cb.into_inner(&mut ctx);
+                let this = ctx.undefined();
+                let args: Vec<Handle<JsValue>> = match result {
+                    Ok(val) => {
+                        let value: Vec<u8>;
+                        let empty_hash = vec![].hash_with_kind(HashKind::Empty);
+                        let current_state_info = if let Some(val) = val {
+                            value = val;
+                            CurrentState::from_bytes(&value)
+                        } else {
+                            CurrentState::new(&empty_hash, BlockHeight(0))
+                        };
+                        let root = JsBuffer::external(&mut ctx, current_state_info.root.to_vec());
+                        let version = ctx.number::<u32>(current_state_info.version.into());
+                        let object = ctx.empty_object();
+                        object.set(&mut ctx, "root", root)?;
+                        object.set(&mut ctx, "version", version)?;
+                        vec![ctx.null().upcast(), object.upcast()]
+                    },
+                    Err(err) => vec![ctx.error(&err)?.upcast()],
+                };
+                callback.call(&mut ctx, this, args)?;
+                Ok(())
+            });
+        })
+    }
 }
 
 impl StateDB {
@@ -409,6 +483,17 @@ impl StateDB {
         Ok(ctx.undefined())
     }
 
+    pub fn js_get_current_state(mut ctx: FunctionContext) -> JsResult<JsUndefined> {
+        let cb = ctx.argument::<JsFunction>(0)?.root(&mut ctx);
+        // Get the `this` value as a `JsBox<Database>`
+        let db = ctx.this().downcast_or_throw::<SharedStateDB, _>(&mut ctx)?;
+        let db = db.borrow();
+        db.get_current_state(cb)
+            .or_else(|err| ctx.throw_error(err.to_string()))?;
+
+        Ok(ctx.undefined())
+    }
+
     pub fn js_exists(mut ctx: FunctionContext) -> JsResult<JsUndefined> {
         let key = ctx.argument::<JsTypedArray<u8>>(0)?.as_slice(&ctx).to_vec();
         let cb = ctx.argument::<JsFunction>(1)?.root(&mut ctx);
@@ -611,3 +696,40 @@ impl StateDB {
         Ok(ctx.undefined())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::{BlockHeight, HashKind, HashWithKind};
+
+    #[test]
+    fn test_current_state_convert() {
+        let block_zero = BlockHeight(0);
+        let block_ten = BlockHeight(10);
+        let block_hundred = BlockHeight(100);
+        let root = Vec::with_capacity(30);
+        let empty_hash = vec![].hash_with_kind(HashKind::Empty);
+        let test_data = vec![
+            (
+                CurrentState::new(&[], block_zero),
+                [vec![], block_zero.to_be_bytes().to_vec()].concat(),
+            ),
+            (
+                CurrentState::new(&[1, 2, 3, 4], block_ten),
+                [vec![1, 2, 3, 4], block_ten.to_be_bytes().to_vec()].concat(),
+            ),
+            (
+                CurrentState::new(&root, block_hundred),
+                [root.clone(), block_hundred.to_be_bytes().to_vec()].concat(),
+            ),
+            (
+                CurrentState::new(&empty_hash, block_zero),
+                [empty_hash.clone(), block_zero.to_be_bytes().to_vec()].concat(),
+            ),
+        ];
+        for (state_as_struct, state_as_bytes) in test_data {
+            assert_eq!(state_as_struct.to_bytes(), state_as_bytes);
+            assert_eq!(CurrentState::from_bytes(&state_as_bytes), state_as_struct);
+        }
+    }
+}

src/types.rs

@@ -20,6 +20,10 @@ pub type ArcMutex<T> = Arc<Mutex<T>>;
 #[derive(Clone, Debug, Copy, PartialEq, Eq)]
 pub struct Height(pub u16);
 
+// Strong type of block height
+#[derive(Clone, Debug, Copy, PartialEq, Eq)]
+pub struct BlockHeight(pub u32);
+
 // Strong type of structure position in Subtree with max value 2 ^ SUBTREE_SIZE
 #[derive(Clone, Debug, Copy, PartialEq, Eq)]
 pub struct StructurePosition(pub u16);
@@ -156,6 +160,34 @@ impl From<f64> for KeyLength {
     }
 }
 
+impl From<KeyLength> for BlockHeight {
+    #[inline]
+    fn from(value: KeyLength) -> BlockHeight {
+        Self(value.0 as u32)
+    }
+}
+
+impl From<u32> for BlockHeight {
+    #[inline]
+    fn from(value: u32) -> BlockHeight {
+        BlockHeight(value)
+    }
+}
+
+impl From<BlockHeight> for u32 {
+    #[inline]
+    fn from(value: BlockHeight) -> u32 {
+        value.0
+    }
+}
+
+impl From<Height> for BlockHeight {
+    #[inline]
+    fn from(value: Height) -> BlockHeight {
+        BlockHeight(value.0 as u32)
+    }
+}
+
 impl Default for SubtreeHeight {
     #[inline]
     fn default() -> Self {
@@ -316,6 +348,13 @@ impl Height {
     }
 }
 
+impl BlockHeight {
+    #[inline]
+    pub fn to_be_bytes(self) -> [u8; 4] {
+        self.0.to_be_bytes()
+    }
+}
+
 impl KeyLength {
     // Cast to u32 and returns with len(4) for JS API
     #[inline]

src/utils.rs

@@ -22,7 +22,7 @@ pub fn get_iteration_mode<'a>(
             .clone()
             .unwrap_or_else(|| vec![255; options.gte.as_ref().unwrap().len()]);
         *opt = if has_prefix {
-            [consts::PREFIX_STATE, lte.as_slice()].concat()
+            [consts::Prefix::STATE, lte.as_slice()].concat()
         } else {
             lte
         };
@@ -33,7 +33,7 @@ pub fn get_iteration_mode<'a>(
             .clone()
             .unwrap_or_else(|| vec![0; options.lte.as_ref().unwrap().len()]);
         *opt = if has_prefix {
-            [consts::PREFIX_STATE, gte.as_slice()].concat()
+            [consts::Prefix::STATE, gte.as_slice()].concat()
         } else {
             gte
         };
@@ -53,7 +53,7 @@ pub fn is_key_out_of_range(
     if options.reverse {
         if let Some(gte) = &options.gte {
             let cmp = if has_prefix {
-                [consts::PREFIX_STATE, gte].concat()
+                [consts::Prefix::STATE, gte].concat()
             } else {
                 gte.to_vec()
             };
@@ -63,7 +63,7 @@ pub fn is_key_out_of_range(
         }
     } else if let Some(lte) = &options.lte {
         let cmp = if has_prefix {
-            [consts::PREFIX_STATE, lte].concat()
+            [consts::Prefix::STATE, lte].concat()
         } else {
             lte.to_vec()
         };