chore: allow consecutive op index for plugin

zhiqiangxu · Aug 14, 2023 · 38718a3 · 38718a3
1 parent 11167f8
commit 38718a3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 23 deletions.
diff --git a/crates/zkwasm/src/runtime/host/external_circuit_plugin.rs b/crates/zkwasm/src/runtime/host/external_circuit_plugin.rs
@@ -71,25 +71,23 @@ impl ModuleImportResolver for ExternalCircuitEnv {
         function_name: &str,
         signature: &wasmi::Signature,
     ) -> Result<wasmi::FuncRef, wasmi::Error> {
-        for (name, function) in &self.functions {
-            if name == function_name {
-                if function.sig.match_wasmi_signature(signature) {
-                    return Ok(FuncInstance::alloc_host(
-                        signature.clone(),
-                        function.op_index,
-                    ));
-                } else {
-                    return Err(wasmi::Error::Instantiation(format!(
-                        "Export `{}` doesnt match expected type {:?}",
-                        function_name, signature
-                    )));
-                }
+        if let Some(function) = self.functions.get(function_name) {
+            if function.sig.match_wasmi_signature(signature) {
+                Ok(FuncInstance::alloc_host(
+                    signature.clone(),
+                    function.op_index,
+                ))
+            } else {
+                Err(wasmi::Error::Instantiation(format!(
+                    "Export `{}` doesnt match expected type {:?}",
+                    function_name, signature
+                )))
             }
+        } else {
+            Err(wasmi::Error::Instantiation(format!(
+                "Export {} not found",
+                function_name
+            )))
         }
-
-        return Err(wasmi::Error::Instantiation(format!(
-            "Export {} not found",
-            function_name
-        )));
     }
 }
diff --git a/crates/zkwasm/src/runtime/host/host_env.rs b/crates/zkwasm/src/runtime/host/host_env.rs
@@ -56,13 +56,10 @@ impl HostEnv {
     pub fn finalize(&mut self) {
         let mut lookup = HashMap::<usize, HostFunction>::new();
 
-        let mut internal_op_allocator_offset = self.external_env.functions.len();
+        let mut internal_op_allocator_offset = 0;
 
         for (name, op) in &self.external_env.functions {
-            assert!(
-                op.op_index < internal_op_allocator_offset,
-                "Specify op index too large."
-            );
+            internal_op_allocator_offset = usize::max(internal_op_allocator_offset, op.op_index);
 
             lookup
                 .insert(
@@ -82,6 +79,8 @@ impl HostEnv {
                 .map(|_| panic!("conflicting op index of foreign function"));
         }
 
+        internal_op_allocator_offset += 1;
+
         for (name, op) in &mut self.internal_env.functions {
             op.index = Some(internal_op_allocator_offset);