fix: sbrk must be called twice (#41)

xcq-api/procedural/src/program/expand/mod.rs

@@ -62,22 +62,19 @@ fn pass_byte_to_host() -> TokenStream2 {
     // TODO check res type to determine the appropriate serializing method
     quote! {
         let res_bytes = res.to_le_bytes();
-        let ptr = polkavm_derive::sbrk(res_bytes.len());
-        if ptr.is_null(){
+        let res_ptr = polkavm_derive::sbrk(0);
+        let end_ptr = polkavm_derive::sbrk(res_bytes.len());
+        if end_ptr.is_null(){
             return 0;
         }
         unsafe {
-            core::ptr::copy_nonoverlapping(res_bytes.as_ptr(),ptr,res_bytes.len());
+            core::ptr::copy_nonoverlapping(res_bytes.as_ptr(),res_ptr,res_bytes.len());
         }
-        (res_bytes.len() as u64) << 32 | (ptr as u64)
+        (res_bytes.len() as u64) << 32 | (res_ptr as u64)
     }
 }
 
 fn generate_main(entrypoint: &EntrypointDef) -> Result<TokenStream2> {
-    let move_to_stack = quote! {
-        let arg_bytes = unsafe {alloc::vec::Vec::from_raw_parts(ptr as *mut u8, size as usize, size as usize)};
-        let mut arg_ptr = arg_bytes.as_ptr() as u32;
-    };
     // Construct call_data
     let mut get_call_data = TokenStream2::new();
     for (arg_type_index, arg_type) in entrypoint.arg_types.iter().enumerate() {
@@ -141,8 +138,7 @@ fn generate_main(entrypoint: &EntrypointDef) -> Result<TokenStream2> {
 
     let main = quote! {
         #[polkavm_derive::polkavm_export]
-        extern "C" fn main(ptr: u32, size:u32) -> u64 {
-            #move_to_stack
+        extern "C" fn main(mut arg_ptr: u32, size:u32) -> u64 {
             #get_call_data
             #call_entrypoint
             #pass_bytes_back

xcq-executor/src/lib.rs

@@ -63,15 +63,20 @@ impl<Ctx: XcqExecutorContext> XcqExecutor<Ctx> {
         let instance_pre = self.linker.instantiate_pre(&module)?;
         let instance = instance_pre.instantiate()?;
 
-        // Args are passed via guest's heap
         let input_ptr = if !input.is_empty() {
-            let ptr = instance
-                .sbrk(input.len() as u32)?
-                .expect("sbrk must be able to allocate memory here");
+            // First sbrk call to get the start of the heap
+            let start_ptr = instance.sbrk(0)?.expect("should not fail because we don't allocate");
+            // Second sbrk call to check the allocation doesn't exceed the heap limit
+            if instance.sbrk(input.len() as u32)?.is_none() {
+                return Err(XcqExecutorError::ExecutionError(polkavm::ExecutionError::Error(
+                    polkavm::Error::from("cannot srk enough memory"),
+                )));
+            };
+            // Args are passed via guest's heap
             instance
-                .write_memory(ptr, input)
+                .write_memory(start_ptr, input)
                 .map_err(|e| XcqExecutorError::ExecutionError(polkavm::ExecutionError::Trap(e)))?;
-            ptr
+            start_ptr
         } else {
             0
         };

xcq-extension/src/lib.rs

@@ -72,7 +72,10 @@ impl<E: ExtensionTuple, P: PermController> XcqExecutorContext for Context<E, P>
                         let res_bytes = E::dispatch(extension_id, &call_bytes)?;
                         tracing::debug!("(host call): res_bytes: {:?}", res_bytes);
                         let res_bytes_len = res_bytes.len();
-                        let res_ptr = caller.sbrk(res_bytes_len as u32).ok_or(ExtensionError::PolkavmError)?;
+                        let res_ptr = caller.sbrk(0).ok_or(ExtensionError::PolkavmError)?;
+                        if caller.sbrk(res_bytes_len as u32).is_none() {
+                            return Err(ExtensionError::PolkavmError);
+                        }
                         caller
                             .write_memory(res_ptr, &res_bytes)
                             .map_err(|_| ExtensionError::PolkavmError)?;

xcq-extension/tests/with_associated_types_works.rs

@@ -122,7 +122,7 @@ fn call_core_works() {
 }
 #[test]
 fn multi_calls_works() {
-    let blob = include_bytes!("../../output/poc-guest-balance-sum-percent.polkavm");
+    let blob = include_bytes!("../../output/poc-guest-sum-balance-percent.polkavm");
     let mut executor = ExtensionsExecutor::<Extensions, ()>::new(InvokeSource::RuntimeAPI);
     let guest = GuestImpl { program: blob.to_vec() };
     let mut input_data = extension_fungibles::EXTENSION_ID.encode();

xcq-test-runner/src/main.rs

@@ -42,11 +42,11 @@ fn main() {
     };
     input_data.extend_from_slice(&method1_encoded);
     input_data.extend_from_slice(&method2.encode());
-    input_data.extend_from_slice(&xcq_extension_fungibles::EXTENSION_ID.encode());
-    let method3 = FungiblesMethod::TotalSupply { asset: 1 };
-    let method3_encoded = method3.encode();
-    input_data.extend_from_slice(&[method3_encoded.len() as u8]);
-    input_data.extend_from_slice(&method3_encoded);
+    // input_data.extend_from_slice(&xcq_extension_fungibles::EXTENSION_ID.encode());
+    // let method3 = FungiblesMethod::TotalSupply { asset: 1 };
+    // let method3_encoded = method3.encode();
+    // input_data.extend_from_slice(&[method3_encoded.len() as u8]);
+    // input_data.extend_from_slice(&method3_encoded);
     tracing::info!("Input data: {:?}", input_data);
     let input = InputImpl {
         method: "main".to_string(),