Paginated datastore-key ABIs (AS path, exploratory)

src/as_execution/abi.rs

@@ -596,6 +596,106 @@ pub(crate) fn assembly_script_get_keys_for(
     })
 }
 
+/// Empty buffer means the bound is unset.
+fn opt_bound(buffer: &[u8]) -> Option<&[u8]> {
+    (!buffer.is_empty()).then_some(buffer)
+}
+
+/// Page size accepted by the paginated datastore-key ABIs: `1..=MAX_DATASTORE_KEYS_PAGE`.
+fn checked_page_count(count: i32) -> ABIResult<u32> {
+    let Ok(count) = u32::try_from(count) else {
+        abi_bail!(format!("negative datastore key page size: {count}"));
+    };
+    if !(1..=crate::MAX_DATASTORE_KEYS_PAGE).contains(&count) {
+        abi_bail!(format!(
+            "datastore key page size must be between 1 and {}, got {}",
+            crate::MAX_DATASTORE_KEYS_PAGE,
+            count
+        ));
+    }
+    Ok(count)
+}
+
+/// Get one page of keys (aka entries) in the datastore.
+///
+/// `prefix` and `start_key` are optional (empty means unset). `start_key` is an
+/// exclusive cursor. `count` must be in `1..=MAX_DATASTORE_KEYS_PAGE`.
+pub(crate) fn assembly_script_get_keys_paginated(
+    mut ctx: FunctionEnvMut<ASEnv>,
+    prefix: i32,
+    start_key: i32,
+    count: i32,
+) -> ABIResult<i32> {
+    abi_with_memory!(ctx, assembly_script_get_keys_paginated, |memory| {
+        let prefix = read_buffer(&memory, &ctx, prefix)?;
+        let start_key = read_buffer(&memory, &ctx, start_key)?;
+        let count = checked_page_count(count)?;
+        let keys = ctx.data().interface.get_keys_paginated(
+            opt_bound(&prefix),
+            opt_bound(&start_key),
+            count,
+        )?;
+        let fmt_keys =
+            ser_bytearray_vec(&keys, keys.len(), crate::MAX_DATASTORE_KEYS_PAGE as usize)?;
+        let ffi_env = ctx.data().get_ffi_env().clone();
+        let ptr = BufferPtr::alloc(&fmt_keys, &ffi_env, &mut ctx)?.offset();
+
+        #[cfg(feature = "execution-trace")]
+        ctx.data_mut().trace.push(AbiTrace {
+            name: "assembly_script_get_keys_paginated".to_string(),
+            params: vec![
+                into_trace_value!(prefix),
+                into_trace_value!(start_key),
+                into_trace_value!(count),
+            ],
+            return_value: fmt_keys.into(),
+            sub_calls: None,
+        });
+        Ok(ptr as i32)
+    })
+}
+
+/// Get one page of keys (aka entries) in the datastore of a given address.
+/// See [`assembly_script_get_keys_paginated`].
+pub(crate) fn assembly_script_get_keys_for_paginated(
+    mut ctx: FunctionEnvMut<ASEnv>,
+    address: i32,
+    prefix: i32,
+    start_key: i32,
+    count: i32,
+) -> ABIResult<i32> {
+    abi_with_memory!(ctx, assembly_script_get_keys_for_paginated, |memory| {
+        let address = read_string(&memory, &ctx, address)?;
+        let prefix = read_buffer(&memory, &ctx, prefix)?;
+        let start_key = read_buffer(&memory, &ctx, start_key)?;
+        let count = checked_page_count(count)?;
+        let keys = ctx.data().interface.get_keys_for_paginated(
+            &address,
+            opt_bound(&prefix),
+            opt_bound(&start_key),
+            count,
+        )?;
+        let fmt_keys =
+            ser_bytearray_vec(&keys, keys.len(), crate::MAX_DATASTORE_KEYS_PAGE as usize)?;
+        let ffi_env = ctx.data().get_ffi_env().clone();
+        let ptr = BufferPtr::alloc(&fmt_keys, &ffi_env, &mut ctx)?.offset();
+
+        #[cfg(feature = "execution-trace")]
+        ctx.data_mut().trace.push(AbiTrace {
+            name: "assembly_script_get_keys_for_paginated".to_string(),
+            params: vec![
+                into_trace_value!(address),
+                into_trace_value!(prefix),
+                into_trace_value!(start_key),
+                into_trace_value!(count),
+            ],
+            return_value: AbiTraceType::ByteArrays(keys.iter().cloned().collect()),
+            sub_calls: None,
+        });
+        Ok(ptr as i32)
+    })
+}
+
 /// sets a key-indexed data entry in the datastore, overwriting existing values
 /// if any
 pub(crate) fn assembly_script_set_data(
@@ -1934,7 +2034,34 @@ pub(crate) fn assembly_script_hash_sha256(
 
 #[cfg(test)]
 mod tests {
-    use crate::as_execution::abi::ser_bytearray_vec;
+    use super::{checked_page_count, opt_bound, ser_bytearray_vec};
+    use crate::{GasCosts, MAX_DATASTORE_KEYS_PAGE};
+
+    #[test]
+    fn test_opt_bound_empty_is_unset() {
+        assert_eq!(opt_bound(&[]), None);
+        assert_eq!(opt_bound(b"ab"), Some(b"ab".as_ref()));
+    }
+
+    #[test]
+    fn test_page_count_bounds() {
+        assert!(checked_page_count(-1).is_err());
+        assert!(checked_page_count(i32::MIN).is_err());
+        assert!(checked_page_count(0).is_err());
+        assert_eq!(checked_page_count(1).unwrap(), 1);
+        assert_eq!(
+            checked_page_count(MAX_DATASTORE_KEYS_PAGE as i32).unwrap(),
+            MAX_DATASTORE_KEYS_PAGE
+        );
+        assert!(checked_page_count(MAX_DATASTORE_KEYS_PAGE as i32 + 1).is_err());
+    }
+
+    #[test]
+    fn test_paginated_ds_keys_gas_is_flat_worst_page() {
+        let costs = GasCosts::default();
+        assert_eq!(costs.assembly_script_get_keys_paginated, 190000);
+        assert_eq!(costs.assembly_script_get_keys_for_paginated, 190000);
+    }
 
     #[test]
     fn test_ser() {

src/as_execution/context.rs

@@ -214,7 +214,7 @@ impl ASContext {
     pub(crate) fn resolver(&self, store: &mut Store) -> (Imports, FunctionEnv<ASEnv>) {
         let fenv = FunctionEnv::new(store, self.env.clone());
 
-        let imports = imports! {
+        let mut imports = imports! {
             "env" => {
                 // Needed by WASM generated by AssemblyScript
                 "abort" =>  Function::new_typed_with_env(store, &fenv, assembly_script_abort),
@@ -289,6 +289,92 @@ impl ASContext {
             },
         };
 
+        // MIP-0002 ABIs. Register them only once the host reports an execution
+        // version at or past their activation, so an updated node rejects the
+        // imports before activation exactly like a non-updated node. A host
+        // that cannot report its version fails closed (old behavior).
+        if self
+            .env
+            .interface
+            .get_interface_version()
+            .is_ok_and(paginated_ds_keys_enabled)
+        {
+            imports.define(
+                "massa",
+                "assembly_script_get_keys_paginated",
+                Function::new_typed_with_env(store, &fenv, assembly_script_get_keys_paginated),
+            );
+            imports.define(
+                "massa",
+                "assembly_script_get_keys_for_paginated",
+                Function::new_typed_with_env(store, &fenv, assembly_script_get_keys_for_paginated),
+            );
+        }
+
         (imports, fenv)
     }
 }
+
+/// Whether the paginated datastore-key ABIs are exposed for a host reporting
+/// this execution component version.
+fn paginated_ds_keys_enabled(interface_version: u32) -> bool {
+    interface_version >= crate::PAGINATED_DS_KEYS_EXECUTION_VERSION
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::{TestInterface, INTERFACE_VERSION};
+    use crate::{CondomLimits, GasCosts, PAGINATED_DS_KEYS_EXECUTION_VERSION};
+    use serial_test::serial;
+    use std::sync::atomic::Ordering;
+    use wasmer::{sys::EngineBuilder, wat2wasm, Module, Store};
+    use wasmer_compiler_singlepass::Singlepass;
+
+    #[test]
+    fn test_paginated_ds_keys_gate_boundaries() {
+        assert!(!paginated_ds_keys_enabled(0));
+        assert!(!paginated_ds_keys_enabled(
+            PAGINATED_DS_KEYS_EXECUTION_VERSION - 1
+        ));
+        assert!(paginated_ds_keys_enabled(
+            PAGINATED_DS_KEYS_EXECUTION_VERSION
+        ));
+        assert!(paginated_ds_keys_enabled(u32::MAX));
+    }
+
+    /// The new imports follow the execution version. The deprecated ones stay
+    /// registered either way, so existing modules still instantiate.
+    #[test]
+    #[serial]
+    fn test_paginated_ds_keys_imports_follow_execution_version() {
+        let engine = EngineBuilder::new(Singlepass::default()).engine();
+        let mut store = Store::new(engine);
+        let module = Module::new(&store, wat2wasm(b"(module)").unwrap()).unwrap();
+        let interface = TestInterface;
+        let ctx = ASContext::new(
+            &interface,
+            module,
+            GasCosts::default(),
+            CondomLimits::default(),
+        );
+
+        INTERFACE_VERSION.store(0, Ordering::SeqCst);
+        let (before, _) = ctx.resolver(&mut store);
+        let before_legacy = before.exists("massa", "assembly_script_get_keys")
+            && before.exists("massa", "assembly_script_get_keys_for");
+        let before_paginated = before.exists("massa", "assembly_script_get_keys_paginated")
+            || before.exists("massa", "assembly_script_get_keys_for_paginated");
+
+        INTERFACE_VERSION.store(PAGINATED_DS_KEYS_EXECUTION_VERSION, Ordering::SeqCst);
+        let (after, _) = ctx.resolver(&mut store);
+        let after_paginated = after.exists("massa", "assembly_script_get_keys_paginated")
+            && after.exists("massa", "assembly_script_get_keys_for_paginated");
+
+        INTERFACE_VERSION.store(0, Ordering::SeqCst);
+
+        assert!(before_legacy);
+        assert!(!before_paginated);
+        assert!(after_paginated);
+    }
+}

src/tests/mod.rs

@@ -13,7 +13,7 @@ pub(crate) static INTERFACE_VERSION: std::sync::atomic::AtomicU32 =
     std::sync::atomic::AtomicU32::new(0);
 
 #[derive(Clone)]
-struct TestInterface;
+pub(crate) struct TestInterface;
 
 impl InterfaceClone for TestInterface {
     fn clone_box(&self) -> Box<dyn Interface> {
@@ -623,6 +623,25 @@ impl Interface for TestInterface {
         todo!()
     }
 
+    fn get_keys_paginated(
+        &self,
+        _prefix: Option<&[u8]>,
+        _start_key: Option<&[u8]>,
+        _count: u32,
+    ) -> Result<BTreeSet<Vec<u8>>> {
+        todo!()
+    }
+
+    fn get_keys_for_paginated(
+        &self,
+        _address: &str,
+        _prefix: Option<&[u8]>,
+        _start_key: Option<&[u8]>,
+        _count: u32,
+    ) -> Result<BTreeSet<Vec<u8>>> {
+        todo!()
+    }
+
     fn raw_get_bytecode(&self) -> Result<Vec<u8>> {
         todo!()
     }

src/types.rs

@@ -273,6 +273,8 @@ pub struct GasCosts {
     pub assembly_script_get_deferred_call_quote: u64,
     pub assembly_script_get_keys: u64,
     pub assembly_script_get_keys_for: u64,
+    pub assembly_script_get_keys_for_paginated: u64,
+    pub assembly_script_get_keys_paginated: u64,
     pub assembly_script_get_op_data: u64,
     pub assembly_script_get_op_keys: u64,
     pub assembly_script_get_op_keys_prefix: u64,
@@ -465,6 +467,10 @@ impl GasCosts {
             ),
             assembly_script_get_keys: get_cost!("assembly_script_get_keys"),
             assembly_script_get_keys_for: get_cost!("assembly_script_get_keys_for"),
+            assembly_script_get_keys_for_paginated: get_cost!(
+                "assembly_script_get_keys_for_paginated"
+            ),
+            assembly_script_get_keys_paginated: get_cost!("assembly_script_get_keys_paginated"),
             assembly_script_get_op_data: get_cost!("assembly_script_get_op_data"),
             assembly_script_get_op_keys: get_cost!("assembly_script_get_op_keys"),
             assembly_script_get_op_keys_prefix: get_cost!("assembly_script_get_op_keys_prefix"),
@@ -621,6 +627,13 @@ impl Default for GasCosts {
             assembly_script_get_deferred_call_quote: 1220,
             assembly_script_get_keys: 1000,
             assembly_script_get_keys_for: 1195,
+            // Flat price for the worst page the cap allows (500 keys):
+            // base + 500 × (per_key + per_key_byte × key length).
+            // 16-byte keys:  2_350 + 500 × 230 ≈ 117k.
+            // 255-byte keys: 117k + 500 × 255 × 0.55 ≈ 187k.
+            // 190k covers every such call, as far as keys go.
+            assembly_script_get_keys_for_paginated: 190000,
+            assembly_script_get_keys_paginated: 190000,
             assembly_script_get_op_data: 50000,
             assembly_script_get_op_keys: 1400,
             assembly_script_get_op_keys_prefix: 1400,
@@ -732,6 +745,16 @@ impl Default for GasCosts {
 /// avoid a massa-versioning dependency.
 pub const WASMV1_RUNTIME_DISABLED_EXECUTION_VERSION: u32 = 2;
 
+/// Execution component version from which the paginated datastore-key ABIs are
+/// exposed to guest modules. Same massa MIP-0002 activation as
+/// [`WASMV1_RUNTIME_DISABLED_EXECUTION_VERSION`]: before it, an updated node
+/// does not resolve the new imports, so instantiation fails exactly as on a
+/// non-updated node. A host that cannot report its version keeps them hidden.
+pub const PAGINATED_DS_KEYS_EXECUTION_VERSION: u32 = WASMV1_RUNTIME_DISABLED_EXECUTION_VERSION;
+
+/// Maximum number of datastore keys one paginated call may return.
+pub const MAX_DATASTORE_KEYS_PAGE: u32 = 500;
+
 #[allow(unused_variables)]
 pub trait Interface: Send + Sync + InterfaceClone {
     fn increment_recursion_counter(&self) -> Result<()>;
@@ -813,14 +836,47 @@ pub trait Interface: Send + Sync + InterfaceClone {
     /// Print function for examples
     fn print(&self, message: &str) -> Result<()>;
 
-    /// Return datastore keys
-    /// Will only return keys with a given prefix if provided in args
+    /// Return datastore keys.
+    /// Only keys with the given prefix are returned when one is provided.
+    ///
+    /// Superseded, from [`PAGINATED_DS_KEYS_EXECUTION_VERSION`], by
+    /// [`Interface::get_keys_paginated`].
     fn get_keys(&self, prefix: Option<&[u8]>) -> Result<BTreeSet<Vec<u8>>>;
 
-    /// Return datastore keys
-    /// Will only return keys with a given prefix if provided in args
+    /// Return datastore keys for an address.
+    /// Only keys with the given prefix are returned when one is provided.
+    ///
+    /// Superseded, from [`PAGINATED_DS_KEYS_EXECUTION_VERSION`], by
+    /// [`Interface::get_keys_for_paginated`].
     fn get_keys_for(&self, address: &str, prefix: Option<&[u8]>) -> Result<BTreeSet<Vec<u8>>>;
 
+    /// Return one page of datastore keys for the current address.
+    ///
+    /// Guest modules can call this only from [`PAGINATED_DS_KEYS_EXECUTION_VERSION`]
+    /// (massa MIP-0002). The AssemblyScript import is omitted before that version.
+    ///
+    /// * `prefix`: only keys with this prefix. `None` matches every key.
+    /// * `start_key`: exclusive resume cursor. `None` starts at the beginning of
+    ///   the range. Pass the last key of the previous page to obtain the next one.
+    /// * `count`: page size, in `1..=`[`MAX_DATASTORE_KEYS_PAGE`]. The ABI rejects
+    ///   any other value. The host must apply the same bound.
+    fn get_keys_paginated(
+        &self,
+        prefix: Option<&[u8]>,
+        start_key: Option<&[u8]>,
+        count: u32,
+    ) -> Result<BTreeSet<Vec<u8>>>;
+
+    /// Return one page of datastore keys for an address.
+    /// See [`Interface::get_keys_paginated`].
+    fn get_keys_for_paginated(
+        &self,
+        address: &str,
+        prefix: Option<&[u8]>,
+        start_key: Option<&[u8]>,
+        count: u32,
+    ) -> Result<BTreeSet<Vec<u8>>>;
+
     fn get_ds_keys_wasmv1(
         &self,
         prefix: &[u8],