Use appropriate relocation type for memory addresses

3 files changed, 74 insertions, 28 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
index 9235dd94..e2ef6c58 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
@@ -203,7 +203,10 @@ public class WasmBinaryFormatInstructions {
                     case I32 -> 0x41;
                     case I64 -> 0x42;
                 });
-                writer.writeMemoryAddressConst(memoryAddressConst.address());
+                writer.writeMemoryAddressConst(
+                    memoryAddressConst.indexType(),
+                    memoryAddressConst.address()
+                );
             }
 
             case WasmTableElementIndexConst tableIndexConst -> {
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
index d731377f..614667da 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
@@ -8,6 +8,7 @@ import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmGlobalName;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmTableIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmTagName;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmIndexType;
 
 import java.io.IOException;
 import java.util.*;
@@ -48,10 +49,18 @@ public class WasmBinaryWriter {
         Leb128Encoding.encodePaddedSignedInt32(value, writer);
     }
 
+    public void writePaddedUnsignedInt64(long value) {
+        Leb128Encoding.encodePaddedUnsignedInt64(value, writer);
+    }
+
     public void writeSignedInt64(long value) {
         Leb128Encoding.encodeSignedInt64(value, writer);
     }
 
+    public void writePaddedSignedInt64(int value) {
+        Leb128Encoding.encodePaddedSignedInt64(value, writer);
+    }
+
     public void writeString(String value) {
         WasmBinaryFormatValues.writeString(value, writer);
     }
@@ -74,9 +83,15 @@ public class WasmBinaryWriter {
         );
     }
 
-    public void writeMemoryAddressConst(WasmMemoryAddress index) throws IOException {
+    public void writeMemoryAddressConst(
+        WasmIndexType indexType,
+        WasmMemoryAddress index
+    ) throws IOException {
         writeRelocatableIndex(
-            WasmRelocationType.MEMORY_ADDR_SLEB,
+            switch (indexType) {
+                case I32 -> WasmRelocationType.MEMORY_ADDR_SLEB;
+                case I64 -> WasmRelocationType.MEMORY_ADDR_SLEB64;
+            },
             OptionalInt.empty(),
             new WasmRelocatableIndex.MemoryAddress(index)
         );
@@ -128,11 +143,7 @@ public class WasmBinaryWriter {
         if (relocatableSectionName.isPresent()) {
             var value = unrelocatedIndex.orElse(0);
             var sectionContentsOffset = writer.position() - sectionContentsStart;
-            if (relocationType.isSigned()) {
-                writePaddedSignedInt32(value);
-            } else {
-                writePaddedUnsignedInt32(value);
-            }
+            writePaddedInt(relocationType, value);
 
             if (!relocationSections.containsKey(sectionIndex)) {
                 relocationSections.put(sectionIndex, WasmRelocationSection.empty(
@@ -148,6 +159,7 @@ public class WasmBinaryWriter {
         } else if (unrelocatedIndex.isEmpty()) {
             throw new IllegalStateException("cannot write relocatable index outside of relocatable section");
         } else {
+            // TODO: handle 64-bit values
             if (relocationType.isSigned()) {
                 writeSignedInt32(unrelocatedIndex.getAsInt());
             } else {
@@ -156,6 +168,29 @@ public class WasmBinaryWriter {
         }
     }
 
+    private void writePaddedInt(WasmRelocationType relocationType, int value) {
+        if (relocationType.isSigned()) {
+            switch (relocationType.indexType()) {
+                case I32 -> {
+                    writePaddedSignedInt32(value);
+                }
+                case I64 -> {
+                    writePaddedSignedInt64(value);
+                }
+            }
+
+        } else {
+            switch (relocationType.indexType()) {
+                case I32 -> {
+                    writePaddedUnsignedInt32(value);
+                }
+                case I64 -> {
+                    writePaddedUnsignedInt64(value);
+                }
+            }
+        }
+    }
+
     public void writeSection(WasmSectionId sectionId, Write writeContents) throws IOException {
         currentSection = Optional.of(sectionId);
         sectionIndex++;
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmRelocationType.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmRelocationType.java
index 1b76028b..c0ea248f 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmRelocationType.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmRelocationType.java
@@ -1,32 +1,36 @@
 package org.shedlang.compiler.backends.wasm.structure.binary;
 
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmIndexType;
+
 public enum WasmRelocationType {
-    FUNCTION_INDEX_LEB(0, false, false),
-    TABLE_INDEX_SLEB(1, true, false),
-    TABLE_INDEX_I32(2, true, false),
-    MEMORY_ADDR_LEB(3, false, true),
-    MEMORY_ADDR_SLEB(4, true, true),
-    MEMORY_ADDR_I32(5, true, true),
-    TYPE_INDEX_LEB(6, false, false),
-    GLOBAL_INDEX_LEB(7, false, false),
-    FUNCTION_OFFSET_I32(8, true, true),
-    SECTION_OFFSET_I32(9, true, true),
-    TAG_INDEX_LEB(10, false, false),
-    GLOBAL_INDEX_I32(13, true, false),
-    MEMORY_ADDR_LEB64(14, false, true),
-    MEMORY_ADDR_SLEB64(15, true, true),
-    MEMORY_ADDR_I64(16, true, true),
-    MEMORY_ADDR_REL_SLEB64(17, true, true),
-    TABLE_INDEX_SLEB64(18, true, false),
-    TABLE_INDEX_I64(19, true, false),
-    TABLE_NUMBER_LEB(20, false, false);
+    FUNCTION_INDEX_LEB(0, WasmIndexType.I32, false, false),
+    TABLE_INDEX_SLEB(1, WasmIndexType.I32, true, false),
+    TABLE_INDEX_I32(2, WasmIndexType.I32, true, false),
+    MEMORY_ADDR_LEB(3, WasmIndexType.I32, false, true),
+    MEMORY_ADDR_SLEB(4, WasmIndexType.I32, true, true),
+    MEMORY_ADDR_I32(5, WasmIndexType.I32, true, true),
+    TYPE_INDEX_LEB(6, WasmIndexType.I32, false, false),
+    GLOBAL_INDEX_LEB(7, WasmIndexType.I32, false, false),
+    FUNCTION_OFFSET_I32(8, WasmIndexType.I32, true, true),
+    SECTION_OFFSET_I32(9, WasmIndexType.I32, true, true),
+    TAG_INDEX_LEB(10, WasmIndexType.I32, false, false),
+    GLOBAL_INDEX_I32(13, WasmIndexType.I32, true, false),
+    MEMORY_ADDR_LEB64(14, WasmIndexType.I64, false, true),
+    MEMORY_ADDR_SLEB64(15, WasmIndexType.I64, true, true),
+    MEMORY_ADDR_I64(16, WasmIndexType.I64, true, true),
+    MEMORY_ADDR_REL_SLEB64(17, WasmIndexType.I64, true, true),
+    TABLE_INDEX_SLEB64(18, WasmIndexType.I64, true, false),
+    TABLE_INDEX_I64(19, WasmIndexType.I64, true, false),
+    TABLE_NUMBER_LEB(20, WasmIndexType.I32, false, false);
 
     private final byte id;
+    private final WasmIndexType indexType;
     private final boolean isSigned;
     private final boolean hasAddend;
 
-    WasmRelocationType(int id, boolean isSigned, boolean hasAddend) {
+    WasmRelocationType(int id, WasmIndexType indexType, boolean isSigned, boolean hasAddend) {
         this.id = (byte) id;
+        this.indexType = indexType;
         this.isSigned = isSigned;
         this.hasAddend = hasAddend;
     }
@@ -35,6 +39,10 @@ public enum WasmRelocationType {
         return id;
     }
 
+    public WasmIndexType indexType() {
+        return indexType;
+    }
+
     public boolean isSigned() {
         return isSigned;
     }