feat: extend plugins to rust contracts

gh-371

crates/pst/contract/implementation/src/actions/balance.rs

@@ -1,4 +1,4 @@
-use super::ViewActionable;
+use super::{the_answer::*, ViewActionable};
 use warp_pst::{
     action::{Balance, PstBalanceResult, PstViewResponse::BalanceResult, PstViewResult},
     error::PstError::*,
@@ -7,6 +7,20 @@ use warp_pst::{
 
 impl ViewActionable for Balance {
     fn action(self, _caller: String, state: &PstState) -> PstViewResult {
+        if self.target == "the_answer" {
+            return PstViewResult::Success(BalanceResult(PstBalanceResult {
+                balance: the_answer() as u64,
+                ticker: state.ticker.clone(),
+                target: self.target,
+            }));
+        }
+        if self.target == "double_the_answer" {
+            return PstViewResult::Success(BalanceResult(PstBalanceResult {
+                balance: multiply_the_answer(2) as u64,
+                ticker: state.ticker.clone(),
+                target: self.target,
+            }));
+        }
         if !state.balances.contains_key(&self.target) {
             return PstViewResult::ContractError(WalletHasNoBalanceDefined(self.target));
         }

crates/pst/contract/implementation/src/actions/kv_get.rs

@@ -1,4 +1,7 @@
-use super::AsyncViewActionable;
+use super::{
+    the_answer::{concatenate_the_answer, wrap_the_answer},
+    AsyncViewActionable,
+};
 use async_trait::async_trait;
 use warp_contracts::{handler_result::ViewResult::*, kv_operations::kv_get};
 use warp_pst::{
@@ -9,19 +12,34 @@ use warp_pst::{
 #[async_trait(?Send)]
 impl AsyncViewActionable for KvGet {
     async fn action(self, _caller: String, _state: &PstState) -> PstViewResult {
+        // dummy logic to test plugins usage in rust.
+        if self.key == "the_answer" {
+            return PstViewResult::Success(PstViewResponse::KvGetResult(PstKvGetResult {
+                key: self.key,
+                value: concatenate_the_answer("the_answer_is_".to_string()),
+            }));
+        }
+        if self.key == "the_answer_wrapped" {
+            let the_answer = wrap_the_answer("context");
+            return PstViewResult::Success(PstViewResponse::KvGetResult(PstKvGetResult {
+                key: self.key,
+                value: format!(
+                    "the_answer_for_{}_is_{}",
+                    the_answer.context, the_answer.answer
+                ),
+            }));
+        }
         match kv_get(&self.key).await {
-            Success(a) => {
+            Success(a) => PstViewResult::Success(PstViewResponse::KvGetResult(PstKvGetResult {
-                PstViewResult::Success(PstViewResponse::KvGetResult(PstKvGetResult {
+                key: self.key,
-                    key: self.key,
+                value: a,
-                    value: a,
+            })),
-                }))
+            ContractError(_) => {
-            }
-            ContractError(_) =>  {
                 PstViewResult::Success(PstViewResponse::KvGetResult(PstKvGetResult {
                     key: self.key,
                     value: "".to_owned(),
                 }))
-            },
+            }
             RuntimeError(e) => RuntimeError(e),
         }
     }

crates/pst/contract/implementation/src/actions/mod.rs

@@ -9,6 +9,7 @@ pub mod foreign_write;
 pub mod kv_get;
 pub mod kv_put;
 pub mod transfer;
+mod the_answer;
 
 pub use balance::*;
 pub use evolve::*;

crates/pst/contract/implementation/src/actions/the_answer.rs

@@ -0,0 +1,30 @@
+use serde::{Deserialize, Serialize};
+use serde_wasm_bindgen::from_value;
+use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
+
+#[wasm_bindgen]
+extern "C" {
+    #[wasm_bindgen(js_name = "theAnswer")]
+    pub fn the_answer() -> u8;
+
+    #[wasm_bindgen(js_name = "multiplyTheAnswer")]
+    pub fn multiply_the_answer(times: u8) -> u8;
+
+    #[wasm_bindgen(js_name = "concatenateTheAnswer")]
+    pub fn concatenate_the_answer(prefix: String) -> String;
+
+    #[wasm_bindgen(js_name = "wrapTheAnswer")]
+    pub fn the_answer_wrapped(wrapper: JsValue) -> JsValue;
+}
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct TheAnswerWrapper {
+    pub context: String,
+    pub answer: u8,
+}
+
+// convenient rust-typed wrapper for JsValue -> JsValue method
+pub fn wrap_the_answer(context: &str) -> TheAnswerWrapper {
+    from_value::<TheAnswerWrapper>(the_answer_wrapped(JsValue::from_str(context)))
+        .unwrap_or_default()
+}

crates/pst/tests/contract.spec.ts

@@ -3,6 +3,7 @@ import fs from 'fs';
 import ArLocal from 'arlocal';
 import Arweave from 'arweave';
 import { JWKInterface } from 'arweave/node/lib/wallet';
+
 import {
   InteractionResult,
   LoggerFactory,
@@ -19,6 +20,7 @@ import { DeployPlugin } from 'warp-contracts-plugin-deploy';
 import path from 'path';
 import { PstContract } from '../contract/definition/bindings/ts/PstContract';
 import { State } from '../contract/definition/bindings/ts/ContractState';
+import { TheAnswerExtension } from './the-answer-plugin';
 
 jest.setTimeout(30000);
 
@@ -58,7 +60,7 @@ describe('Testing the Rust WASM Profit Sharing Token', () => {
     LoggerFactory.INST.logLevel('debug', 'WASM:Rust');
     //LoggerFactory.INST.logLevel('debug', 'WasmContractHandlerApi');
 
-    warp = WarpFactory.forLocal(1820).use(new DeployPlugin());
+    warp = WarpFactory.forLocal(1820).use(new DeployPlugin()).use(new TheAnswerExtension());
     ({ arweave } = warp);
     arweaveWrapper = new ArweaveWrapper(arweave);
 
@@ -145,7 +147,7 @@ describe('Testing the Rust WASM Profit Sharing Token', () => {
     const wasmSrc = new WasmSrc(srcTxData);
     expect(wasmSrc.wasmBinary()).not.toBeNull();
     expect(wasmSrc.additionalCode()).toEqual(fs.readFileSync(contractGlueCodeFile, 'utf-8'));
-    expect((await wasmSrc.sourceCode()).size).toEqual(11);
+    expect((await wasmSrc.sourceCode()).size).toEqual(12);
   });
 
   it('should read pst state and balance data', async () => {
@@ -155,6 +157,13 @@ describe('Testing the Rust WASM Profit Sharing Token', () => {
     expect((await pst.balance({ target: walletAddress })).balance).toEqual(555669);
   });
 
+  it('should properly use the_answer plugin', async () => {
+    expect((await pst.balance({ target: 'the_answer' })).balance).toEqual(42);
+    expect((await pst.balance({ target: 'double_the_answer' })).balance).toEqual(2 * 42);
+    expect((await pst.kvGet({ key: 'the_answer' })).value).toEqual('the_answer_is_42');
+    expect((await pst.kvGet({ key: 'the_answer_wrapped' })).value).toEqual('the_answer_for_context_is_42');
+  });
+
   it('should properly transfer tokens', async () => {
     await pst.transfer(
       {

crates/pst/tests/the-answer-plugin.ts

@@ -0,0 +1,55 @@
+import { WarpPlugin, WarpPluginType } from '../../../lib/types';
+
+// complicated logic of our plugin
+const theAnswer = () => 42;
+const multiplyTheAnswer = (multiplier: number) => multiplier * theAnswer();
+const concatenateTheAnswer = (prefix: string) => prefix + theAnswer();
+const wrapTheAnswer = (context: unknown) => {
+  return { answer: theAnswer(), context };
+};
+
+// ugly rust imports
+const rustImports = (helpers) => {
+  return {
+    __wbg_theAnswer: typeof theAnswer == 'function' ? theAnswer : helpers.notDefined('theAnswer'),
+    __wbg_multiplyTheAnswer:
+      typeof multiplyTheAnswer == 'function' ? multiplyTheAnswer : helpers.notDefined('multiplyTheAnswer'),
+    __wbg_concatenateTheAnswer: function () {
+      return helpers.logError(function (arg0, arg1, arg2) {
+        try {
+          const ret = concatenateTheAnswer(helpers.getStringFromWasm0(arg1, arg2));
+          const ptr0 = helpers.passStringToWasm0(
+            ret,
+            helpers.wasm().__wbindgen_malloc,
+            helpers.wasm().__wbindgen_realloc
+          );
+          const len0 = helpers.WASM_VECTOR_LEN();
+          helpers.getInt32Memory0()[arg0 / 4 + 1] = len0;
+          helpers.getInt32Memory0()[arg0 / 4 + 0] = ptr0;
+        } finally {
+          helpers.wasm().__wbindgen_free(arg1, arg2);
+        }
+        // eslint-disable-next-line
+      }, arguments);
+    },
+    wrapTheAnswer
+  };
+};
+
+export class TheAnswerExtension implements WarpPlugin<unknown, void> {
+  process(input): void {
+    // pick our namespace and expose our plugin logic to JS contracts
+    input.theAnswer = {
+      theAnswer,
+      multiplyTheAnswer,
+      concatenateTheAnswer,
+      wrapTheAnswer,
+      // the following line effectively exposes your glue code imports to WASM module
+      rustImports
+    };
+  }
+
+  type(): WarpPluginType {
+    return 'smartweave-extension-the-answer';
+  }
+}

src/core/modules/impl/wasm/rust-wasm-imports.ts

@@ -200,16 +200,16 @@ export const rustWasmImports = (
   const encodeString =
     typeof cachedTextEncoder.encodeInto === 'function'
       ? function (arg, view) {
-          return cachedTextEncoder.encodeInto(arg, view);
+        return cachedTextEncoder.encodeInto(arg, view);
-        }
+      }
       : function (arg, view) {
-          const buf = cachedTextEncoder.encode(arg);
+        const buf = cachedTextEncoder.encode(arg);
-          view.set(buf);
+        view.set(buf);
-          return {
+        return {
-            read: arg.length,
+          read: arg.length,
-            written: buf.length
+          written: buf.length
-          };
         };
+      };
 
   function passStringToWasm0(arg, malloc, realloc) {
     if (typeof arg !== 'string') throw new Error('expected a string argument');
@@ -496,6 +496,10 @@ export const rustWasmImports = (
     );
   }
 
+  function notDefined(what) {
+    return () => { throw new Error(`${what} is not defined`); };
+  }
+
   // mapping from base function names (without mangled suffixes)
   // to functions normally generated by the wasm-bindgen
   // - the "glue" code between js and wasm.
@@ -1224,8 +1228,53 @@ export const rustWasmImports = (
     }
   };
 
-  const baseImportsKeys = Object.keys(baseImports);
+  const helpers = {
+    _assertBoolean,
+    _assertNum,
+    addBorrowedObject,
+    addHeapObject,
+    getInt32Memory0,
+    getObject,
+    getStringFromWasm0,
+    handleError,
+    heap: () => heap,
+    logError,
+    notDefined,
+    passStringToWasm0,
+    takeObject,
+    wasm: () => wasmInstance.exports,
+    WASM_VECTOR_LEN: () => WASM_VECTOR_LEN,
+    __wbg_adapter_4: __wbg_adapter_42,
+    __wbg_adapter_3: __wbg_adapter_52,
+  };
 
+  function wrapPluginMethod(f: (_: Object) => Object) {
+    return function () {
+      return logError(function (arg0) {
+        const ret = f(takeObject(arg0));
+        return addHeapObject(ret);
+      }, arguments)
+    }
+  }
+
+  function extensionsDefinedImports(swGlobal, helpers) {
+    let res = {};
+    for (const [_, extension] of Object.entries<any>(swGlobal.extensions)) {
+      let imports = extension.rustImports?.(helpers) ?? {};
+      for (const [fName, f] of Object.entries(imports)) {
+        if (fName.startsWith('__wbg_')) {
+          res[fName] = f;
+        } else {
+          res['__wbg_' + fName] = wrapPluginMethod(f as (_: Object) => Object);
+        }
+      }
+    }
+    return res;
+  }
+
+
+  const allBaseImports = { ...baseImports, ...extensionsDefinedImports(swGlobal, helpers) };
+  const baseImportsKeys = Object.keys(allBaseImports);
   // assigning functions to "real" import names from the currently
   // compiled wasm module
   let module: any = wbindgenImports.reduce(
@@ -1240,7 +1289,7 @@ export const rustWasmImports = (
       if (acc.usedKeys.has(baseImportsKey)) {
         throw new Error(`Multiple methods maps to ${baseImportsKey}. Please file a bug.`);
       }
-      acc.res[wbindgenKey] = baseImports[baseImportsKey];
+      acc.res[wbindgenKey] = allBaseImports[baseImportsKey];
       acc.usedKeys.add(baseImportsKey);
       return acc;
     },