Drop MemoryRegistry

Author: jirkamarsik

wasm/src/org.graalvm.wasm.test/src/org/graalvm/wasm/test/WasmFileSuite.java

@@ -76,7 +76,6 @@
 import org.graalvm.polyglot.Value;
 import org.graalvm.polyglot.io.IOAccess;
 import org.graalvm.wasm.GlobalRegistry;
-import org.graalvm.wasm.MemoryRegistry;
 import org.graalvm.wasm.RuntimeState;
 import org.graalvm.wasm.WasmContext;
 import org.graalvm.wasm.WasmFunctionInstance;
@@ -216,7 +215,6 @@ private void runInContext(WasmCase testCase, Context context, List<Source> sourc
             }
 
             final WasmContext wasmContext = WasmContext.get(null);
-            final List<WasmInstance> instanceList = moduleInstances.stream().map(i -> toWasmInstance(i)).toList();
 
             final Value testFunction;
             final String entryPoint = testCase.options().getProperty("entry-point");
@@ -261,6 +259,7 @@ private void runInContext(WasmCase testCase, Context context, List<Source> sourc
                 } finally {
                     // Context may have already been closed, e.g. by __wasi_proc_exit.
                     if (!wasmContext.environment().getContext().isClosed()) {
+                        Collection<WasmInstance> instanceList = wasmContext.contextStore().moduleInstances().values();
                         // Save context state, and check that it's consistent with the previous one.
                         if (iterationNeedsStateCheck(i)) {
                             final ContextState contextState = saveContext(wasmContext, instanceList);
@@ -275,8 +274,8 @@ private void runInContext(WasmCase testCase, Context context, List<Source> sourc
                         final boolean reinitMemory = requiresZeroMemory || iterationNeedsStateCheck(i + 1);
                         if (reinitMemory) {
                             for (WasmInstance instance : instanceList) {
-                                for (int j = 0; j < instance.store().memories().count(); ++j) {
-                                    WasmMemoryLibrary.getUncached().reset(instance.store().memories().memory(j));
+                                for (int j = 0; j < instance.module().memoryCount(); ++j) {
+                                    WasmMemoryLibrary.getUncached().reset(instance.memory(j));
                                 }
                                 for (int j = 0; j < instance.module().tableCount(); ++j) {
                                     instance.table(j).reset();
@@ -637,25 +636,28 @@ protected String suiteName() {
     }
 
     private static InstanceState saveInstanceState(WasmInstance instance) {
-        final MemoryRegistry memories = instance.store().memories().duplicate();
+        final WasmMemory[] memories = new WasmMemory[instance.module().memoryCount()];
+        for (int i = 0; i < memories.length; i++) {
+            memories[i] = WasmMemoryLibrary.getUncached().duplicate(instance.memory(i));
+        }
         final GlobalRegistry globals = instance.globals().duplicate();
         return new InstanceState(memories, globals);
     }
 
-    private static ContextState saveContext(WasmContext context, List<WasmInstance> instances) {
+    private static ContextState saveContext(WasmContext context, Collection<WasmInstance> instances) {
         return new ContextState(
                         instances.stream().map(instance -> saveInstanceState(instance)).toList(),
                         context.fdManager().size());
     }
 
     private static void assertInstanceEqual(InstanceState expectedState, InstanceState actualState) {
         // Compare memories
-        final MemoryRegistry expectedMemories = expectedState.memories();
-        final MemoryRegistry actualMemories = actualState.memories();
-        Assert.assertEquals("Mismatch in memory counts.", expectedMemories.count(), actualMemories.count());
-        for (int i = 0; i < expectedMemories.count(); i++) {
-            final WasmMemory expectedMemory = expectedMemories.memory(i);
-            final WasmMemory actualMemory = actualMemories.memory(i);
+        final WasmMemory[] expectedMemories = expectedState.memories();
+        final WasmMemory[] actualMemories = actualState.memories();
+        Assert.assertEquals("Mismatch in memory counts.", expectedMemories.length, actualMemories.length);
+        for (int i = 0; i < expectedMemories.length; i++) {
+            final WasmMemory expectedMemory = expectedMemories[i];
+            final WasmMemory actualMemory = actualMemories[i];
             if (expectedMemory == null) {
                 Assert.assertNull("Memory should be null", actualMemory);
             } else {
@@ -690,7 +692,7 @@ private static void assertContextEqual(ContextState expectedState, ContextState
         Assert.assertEquals("Mismatch in file descriptor counts.", expectedState.openedFdCount, actualState.openedFdCount);
     }
 
-    private record InstanceState(MemoryRegistry memories, GlobalRegistry globals) {
+    private record InstanceState(WasmMemory[] memories, GlobalRegistry globals) {
     }
 
     private record ContextState(List<InstanceState> instanceState, int openedFdCount) {

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/Linker.java

@@ -481,11 +481,8 @@ void resolveMemoryImport(WasmStore store, WasmInstance instance, ImportDescripto
         final String importedModuleName = importDescriptor.moduleName();
         final String importedMemoryName = importDescriptor.memberName();
         final Runnable resolveAction = () -> {
-            final WasmMemory importedMemory;
-            final WasmMemory externalMemory = lookupImportObject(instance, importDescriptor, imports, WasmMemory.class);
-            if (externalMemory != null) {
-                final int contextMemoryIndex = store.memories().register(externalMemory);
-                importedMemory = store.memories().memory(contextMemoryIndex);
+            WasmMemory importedMemory = lookupImportObject(instance, importDescriptor, imports, WasmMemory.class);
+            if (importedMemory != null) {
                 assert memoryIndex == importDescriptor.targetIndex();
             } else {
                 // WASIp1 memory import should have been resolved via ImportValueSupplier above.

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/MemoryRegistry.java

@@ -1534,9 +1534,7 @@ public void allocateMemory(int index, long declaredMinSize, long declaredMaxSize
             } else {
                 wasmMemory = WasmMemoryFactory.createMemory(declaredMinSize, declaredMaxSize, indexType64, shared, useUnsafeMemory, directByteBufferMemoryAccess, context);
             }
-            final int memoryAddress = store.memories().register(wasmMemory);
-            final WasmMemory allocatedMemory = store.memories().memory(memoryAddress);
-            instance.setMemory(index, allocatedMemory);
+            instance.setMemory(index, wasmMemory);
         });
     }
 

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/SymbolTable.java

@@ -186,9 +186,7 @@ static List<LinkAction> recreateLinkActions(WasmModule module) {
                     limits.checkMemoryInstanceSize(memoryMinSize, memoryIndexType64);
                     final WasmMemory wasmMemory = WasmMemoryFactory.createMemory(memoryMinSize, memoryMaxSize, memoryIndexType64, memoryShared,
                                     context.getContextOptions().useUnsafeMemory(), context.getContextOptions().directByteBufferMemoryAccess(), context);
-                    final int address = store.memories().register(wasmMemory);
-                    final WasmMemory allocatedMemory = store.memories().memory(address);
-                    instance.setMemory(memoryIndex, allocatedMemory);
+                    instance.setMemory(memoryIndex, wasmMemory);
                 });
             }
         }

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/WasmInstantiator.java

@@ -68,7 +68,6 @@
 public final class WasmStore implements TruffleObject {
     private final WasmContext context;
     private final WasmLanguage language;
-    private final MemoryRegistry memoryRegistry;
     private final Linker linker;
     private final Map<String, WasmInstance> moduleInstances;
     private final FdManager filesManager;
@@ -78,7 +77,6 @@ public WasmStore(WasmContext context, WasmLanguage language) {
         this.context = context;
         this.language = language;
         this.contextOptions = context.getContextOptions();
-        this.memoryRegistry = new MemoryRegistry();
         this.moduleInstances = new LinkedHashMap<>();
         this.linker = new Linker();
         this.filesManager = context.fdManager();
@@ -96,10 +94,6 @@ public WasmLanguage language() {
         return language;
     }
 
-    public MemoryRegistry memories() {
-        return memoryRegistry;
-    }
-
     public Linker linker() {
         return linker;
     }