diff --git a/packages/wasm-privacy-coin/pom.xml b/packages/wasm-privacy-coin/pom.xml index 0d4f7cd3550..e16a26d55fc 100644 --- a/packages/wasm-privacy-coin/pom.xml +++ b/packages/wasm-privacy-coin/pom.xml @@ -20,11 +20,36 @@ - + com.dylibso.chicory runtime - 1.1.1 + 1.7.5 + + + + + com.dylibso.chicory + compiler + 1.7.5 + + + + + org.ow2.asm + asm + 9.9.1 + + + org.ow2.asm + asm-commons + 9.9.1 + + + org.ow2.asm + asm-util + 9.9.1 diff --git a/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java b/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java index 8017c574195..2a53bc4bdbc 100644 --- a/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java +++ b/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java @@ -45,7 +45,7 @@ private static ShieldedMerkleTree create(Consumer init) { ok = true; return new ShieldedMerkleTree(bridge); } finally { - if (!ok) bridge.close(); + if (!ok) bridge.discard(); } } @@ -193,11 +193,7 @@ private static void requireU32(long value, String name) { @Override public void close() { - try { - bridge.call("drop_tree"); - } catch (Exception ignored) { - // Best-effort: drop the in-WASM tree before the instance goes away. - } + // WasmBridge.close() drops the in-WASM tree and returns the Instance to the pool. bridge.close(); } } diff --git a/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/WasmBridge.java b/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/WasmBridge.java index f60fb57713b..2b5f130bc40 100644 --- a/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/WasmBridge.java +++ b/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/WasmBridge.java @@ -3,14 +3,19 @@ import com.bitgo.wasm.privacycoin.WasmException; import com.bitgo.wasm.privacycoin.proto.Response; import com.bitgo.wasm.privacycoin.proto.WasmError; +import com.dylibso.chicory.compiler.MachineFactoryCompiler; import com.dylibso.chicory.runtime.ExportFunction; import com.dylibso.chicory.runtime.Instance; +import com.dylibso.chicory.runtime.Machine; import com.dylibso.chicory.wasm.Parser; +import com.dylibso.chicory.wasm.WasmModule; import com.google.protobuf.InvalidProtocolBufferException; import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.InputStream; +import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.function.Function; /** * Low-level bridge to the shielded-tree WASM instance. @@ -22,10 +27,41 @@ * Each call writes a proto-encoded request into WASM linear memory, invokes the named * export, then reads the {@link Response} proto from the LAST_RESULT buffer. * + *

Performance: The parsed {@link WasmModule}, JIT-compiled machine factory, and + * instantiated {@link Instance} objects are all cached/pooled. The module is parsed once per + * JVM (it is immutable and safe to share). The WASM functions are JIT-compiled to JVM bytecode + * once via {@link MachineFactoryCompiler} — this replaces Chicory's default interpreter with + * native-speed bytecode execution. Instances are pooled: on {@link #close()}, the in-WASM tree + * is dropped via the {@code drop_tree} export and the Instance is returned to a pool for reuse + * by the next {@code WasmBridge}. + * *

Thread safety: Not thread-safe. One bridge per {@link ShieldedMerkleTree}. + * The cached module and instance pool are safe for concurrent access. */ final class WasmBridge implements AutoCloseable { + /** + * Lazily-parsed, JVM-wide cached WASM module. {@link WasmModule} is immutable after + * construction, so it is safe to share across threads and to build multiple independent + * {@link Instance}s from it. + */ + private static volatile WasmModule cachedModule; + + /** + * JIT-compiled machine factory cached alongside the module. The {@link MachineFactoryCompiler} + * translates WASM functions into JVM bytecode at load time; the resulting factory is reused + * by every {@link Instance}, so the compilation cost is paid once per JVM. This replaces + * Chicory's default interpreter with native-speed bytecode execution. + */ + private static volatile Function cachedMachineFactory; + + /** + * Pool of previously-used {@link Instance} objects whose in-WASM tree has been dropped. + * Reusing an Instance avoids the instantiation cost; the {@code from_frontier} / + * {@code from_state} export re-initializes the tree state in the existing linear memory. + */ + private static final ConcurrentLinkedDeque instancePool = new ConcurrentLinkedDeque<>(); + private final Instance instance; private final ExportFunction fnAlloc; private final ExportFunction fnDealloc; @@ -33,28 +69,77 @@ final class WasmBridge implements AutoCloseable { private final ExportFunction fnResultLen; WasmBridge() { - byte[] wasmBytes; + Instance inst = instancePool.pollFirst(); + if (inst == null) { + try { + inst = Instance.builder(loadModule()) + .withMachineFactory(loadMachineFactory()) + .build(); + } catch (Exception e) { + throw new IllegalStateException("Failed to instantiate WASM module", e); + } + } + this.instance = inst; + + this.fnAlloc = instance.export("alloc"); + this.fnDealloc = instance.export("dealloc"); + this.fnResultPtr = instance.export("last_result_ptr"); + this.fnResultLen = instance.export("last_result_len"); + } + + /** + * Returns the cached parsed module, parsing it on first use. Subsequent calls reuse + * the same immutable {@link WasmModule}, eliminating redundant parse overhead when + * multiple {@link ShieldedMerkleTree} instances are created in the same JVM (e.g. + * in tests or when re-initializing after a reorg). + */ + private static WasmModule loadModule() { + WasmModule module = cachedModule; + if (module != null) { + return module; + } + synchronized (WasmBridge.class) { + module = cachedModule; + if (module == null) { + byte[] wasmBytes = readWasmBytes(); + module = Parser.parse(new ByteArrayInputStream(wasmBytes)); + cachedModule = module; + } + return module; + } + } + + /** + * Returns the cached JIT-compiled machine factory, compiling it on first use. The + * {@link MachineFactoryCompiler} translates all WASM functions to JVM bytecode once; + * the resulting {@link Machine} is then used by every {@link Instance} built from + * the same module, providing native-speed execution instead of interpretation. + */ + private static Function loadMachineFactory() { + Function factory = cachedMachineFactory; + if (factory != null) { + return factory; + } + synchronized (WasmBridge.class) { + factory = cachedMachineFactory; + if (factory == null) { + factory = MachineFactoryCompiler.compile(loadModule()); + cachedMachineFactory = factory; + } + return factory; + } + } + + private static byte[] readWasmBytes() { try (InputStream is = WasmBridge.class.getResourceAsStream("/wasm/privacy_coin.wasm")) { if (is == null) { throw new IllegalStateException( "WASM binary not found on classpath: /wasm/privacy_coin.wasm"); } - wasmBytes = is.readAllBytes(); + return is.readAllBytes(); } catch (IOException e) { throw new IllegalStateException("Failed to load WASM binary", e); } - - try { - var module = Parser.parse(new ByteArrayInputStream(wasmBytes)); - this.instance = Instance.builder(module).build(); - } catch (Exception e) { - throw new IllegalStateException("Failed to instantiate WASM module", e); - } - - this.fnAlloc = instance.export("alloc"); - this.fnDealloc = instance.export("dealloc"); - this.fnResultPtr = instance.export("last_result_ptr"); - this.fnResultLen = instance.export("last_result_len"); } private final class WasmBuffer implements AutoCloseable { @@ -114,8 +199,27 @@ static WasmException toWasmException(WasmError error) { return new WasmException(error.getCode(), error.getMessage()); } + /** + * Drops the in-WASM tree and returns the {@link Instance} to the pool for reuse. + * If {@code drop_tree} fails (e.g. the tree was never initialized), the Instance is + * still pooled — the next {@code from_frontier} / {@code from_state} call will + * re-initialize it regardless. + */ @Override public void close() { - // Chicory Instance has no explicit close method; nothing to do here. + try { + instance.export("drop_tree").apply(); + } catch (Exception ignored) { + // Best-effort cleanup; the instance is still reusable. + } + instancePool.addFirst(instance); + } + + /** + * Discards this bridge without returning the {@link Instance} to the pool. Use this + * when initialization failed and the WASM linear memory may be in an inconsistent state. + */ + void discard() { + // Instance is simply abandoned; Chicory has no explicit close. } }