wasm gc: fix issues in disassembler

core/src/main/java/org/teavm/backend/wasm/disasm/Disassembler.java

@@ -15,6 +15,7 @@
  */
 package org.teavm.backend.wasm.disasm;
 
+import com.carrotsearch.hppc.IntArrayList;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
@@ -23,6 +24,8 @@ import java.nio.file.Files;
 import java.util.function.Consumer;
 import org.teavm.backend.wasm.parser.AddressListener;
 import org.teavm.backend.wasm.parser.CodeSectionParser;
+import org.teavm.backend.wasm.parser.FunctionSectionListener;
+import org.teavm.backend.wasm.parser.FunctionSectionParser;
 import org.teavm.backend.wasm.parser.GlobalSectionParser;
 import org.teavm.backend.wasm.parser.ImportSectionListener;
 import org.teavm.backend.wasm.parser.ImportSectionParser;
@@ -30,11 +33,14 @@ import org.teavm.backend.wasm.parser.ModuleParser;
 import org.teavm.backend.wasm.parser.NameSectionListener;
 import org.teavm.backend.wasm.parser.NameSectionParser;
 import org.teavm.backend.wasm.parser.TypeSectionParser;
+import org.teavm.backend.wasm.parser.WasmHollowFunctionType;
 import org.teavm.common.AsyncInputStream;
 import org.teavm.common.ByteArrayAsyncInputStream;
 
 public final class Disassembler {
     private DisassemblyWriter writer;
+    private WasmHollowFunctionType[] functionTypes;
+    private int[] functionTypeRefs;
 
     public Disassembler(DisassemblyWriter writer) {
         this.writer = writer;
@@ -104,6 +110,7 @@ public final class Disassembler {
                 writer.setAddressOffset(pos);
                 var sectionParser = new TypeSectionParser(typeWriter);
                 sectionParser.parse(writer.addressListener, bytes);
+                functionTypes = typeWriter.getFunctionTypes();
                 writer.flush();
             };
         } else if (code == 2) {
@@ -111,19 +118,35 @@ public final class Disassembler {
                 var parser = new ImportSectionParser(importListener);
                 parser.parse(AddressListener.EMPTY, bytes);
             };
+        } else if (code == 3) {
+            return bytes -> {
+                var signatures = new IntArrayList();
+                for (var i = 0; i < importListener.count; ++i) {
+                    signatures.add(0);
+                }
+                var parser = new FunctionSectionParser(new FunctionSectionListener() {
+                    @Override
+                    public void function(int index, int typeIndex) {
+                        signatures.add(typeIndex);
+                    }
+                });
+                parser.parse(AddressListener.EMPTY, bytes);
+                functionTypeRefs = signatures.toArray();
+            };
         } else if (code == 6) {
             return bytes -> {
                 writer.write("(; global section size: " + bytes.length + " ;)").eol();
                 var globalWriter = new DisassemblyGlobalSectionListener(writer, nameProvider);
                 writer.setAddressOffset(pos);
                 var sectionParser = new GlobalSectionParser(globalWriter);
-                sectionParser.setFunctionIndexOffset(importListener.count);
                 sectionParser.parse(writer.addressListener, bytes);
                 writer.flush();
             };
         } else if (code == 10) {
             return bytes -> {
                 var disassembler = new DisassemblyCodeSectionListener(writer, nameProvider);
+                disassembler.setFunctionTypes(functionTypes);
+                disassembler.setFunctionTypeRefs(functionTypeRefs);
                 writer.setAddressOffset(pos);
                 writer.write("(; code section size: " + bytes.length + " ;)").eol();
                 var sectionParser = new CodeSectionParser(disassembler);

core/src/main/java/org/teavm/backend/wasm/disasm/DisassemblyCodeSectionListener.java

@@ -17,17 +17,29 @@ package org.teavm.backend.wasm.disasm;
 
 import org.teavm.backend.wasm.parser.CodeListener;
 import org.teavm.backend.wasm.parser.CodeSectionListener;
+import org.teavm.backend.wasm.parser.WasmHollowFunctionType;
 import org.teavm.backend.wasm.parser.WasmHollowType;
 
 public class DisassemblyCodeSectionListener extends BaseDisassemblyListener implements CodeSectionListener {
     private int currentFunctionId;
+    private int currentFunctionParameterCount;
     private DisassemblyCodeListener codeListener;
+    private WasmHollowFunctionType[] functionTypes;
+    private int[] functionTypeRefs;
 
     public DisassemblyCodeSectionListener(DisassemblyWriter writer, NameProvider nameProvider) {
         super(writer, nameProvider);
         codeListener = new DisassemblyCodeListener(writer, nameProvider);
     }
 
+    public void setFunctionTypes(WasmHollowFunctionType[] functionTypes) {
+        this.functionTypes = functionTypes;
+    }
+
+    public void setFunctionTypeRefs(int[] functionTypeRefs) {
+        this.functionTypeRefs = functionTypeRefs;
+    }
+
     @Override
     public boolean functionStart(int index, int size) {
         currentFunctionId = index;
@@ -37,7 +49,38 @@ public class DisassemblyCodeSectionListener extends BaseDisassemblyListener impl
         if (name != null) {
             writer.write(" $").write(name);
         }
-        writer.endLinkTarget().indent().eol();
+        writer.endLinkTarget();
+
+        var typeRef = functionTypeRefs[index];
+        writer.write(" (type ");
+        writeTypeRef(typeRef);
+        writer.write(")");
+        writer.indent().eol();
+
+        var type = typeRef < functionTypes.length ? functionTypes[typeRef] : null;
+        if (type != null) {
+            currentFunctionParameterCount = type.parameterTypes.length;
+            for (var i = 0; i < type.parameterTypes.length; ++i) {
+                writer.write("(param ");
+                writer.startLinkTarget("l" + index + "." + i).write(" (; " + i + " ;)");
+                var paramName = nameProvider.local(index, i);
+                if (paramName != null) {
+                    writer.write(" $").write(paramName);
+                }
+                writer.endLinkTarget();
+                writer.write(" ");
+                writeType(type.parameterTypes[i]);
+                writer.write(")").eol();
+            }
+            for (var i = 0; i < type.returnTypes.length; ++i) {
+                writer.write("(result ");
+                writeType(type.returnTypes[i]);
+                writer.write(")").eol();
+            }
+        } else {
+            currentFunctionParameterCount = 0;
+        }
+
         return true;
     }
 
@@ -51,7 +94,7 @@ public class DisassemblyCodeSectionListener extends BaseDisassemblyListener impl
         writer.address();
         for (int i = 0; i < count; ++i) {
             writer.write("(local ");
-            var id = i + start;
+            var id = i + start + currentFunctionParameterCount;
             writer.startLinkTarget("l" + currentFunctionId + "." + id).write("(; " + id + " ;)");
             var name = nameProvider.local(currentFunctionId, id);
             if (name != null) {

core/src/main/java/org/teavm/backend/wasm/disasm/DisassemblyTypeSectionListener.java

@@ -15,7 +15,10 @@
  */
 package org.teavm.backend.wasm.disasm;
 
+import java.util.ArrayList;
+import java.util.List;
 import org.teavm.backend.wasm.parser.TypeSectionListener;
+import org.teavm.backend.wasm.parser.WasmHollowFunctionType;
 import org.teavm.backend.wasm.parser.WasmHollowStorageType;
 import org.teavm.backend.wasm.parser.WasmHollowType;
 
@@ -25,6 +28,9 @@ public class DisassemblyTypeSectionListener extends BaseDisassemblyListener impl
     private int currentTypeIndex;
     private int fieldIndex;
     private boolean needsFieldIndex;
+    private List<WasmHollowFunctionType> functionTypes = new ArrayList<>();
+    private List<WasmHollowType> parameterTypes = new ArrayList<>();
+    private List<WasmHollowType> resultTypes = new ArrayList<>();
 
     public DisassemblyTypeSectionListener(DisassemblyWriter writer, NameProvider nameProvider) {
         super(writer, nameProvider);
@@ -42,6 +48,7 @@ public class DisassemblyTypeSectionListener extends BaseDisassemblyListener impl
 
     @Override
     public void startType(int index, boolean open, int[] supertypes) {
+        functionTypes.add(null);
         currentTypeIndex = index;
         writer.address().write("(type ");
         writer.startLinkTarget("t" + index).write("(; ").write(String.valueOf(index)).write(" ;) ");
@@ -119,6 +126,11 @@ public class DisassemblyTypeSectionListener extends BaseDisassemblyListener impl
 
     @Override
     public void resultType(WasmHollowType type) {
+        if (emittingReturn) {
+            resultTypes.add(type);
+        } else {
+            parameterTypes.add(type);
+        }
         writer.address().write("(").write(emittingReturn ? "result" : "param").write(" ");
         writeType(type);
         writer.write(")").eol();
@@ -126,6 +138,11 @@ public class DisassemblyTypeSectionListener extends BaseDisassemblyListener impl
 
     @Override
     public void endFuncType() {
+        var type = new WasmHollowFunctionType(parameterTypes.toArray(new WasmHollowType[0]),
+                resultTypes.toArray(new WasmHollowType[0]));
+        functionTypes.set(currentTypeIndex, type);
+        parameterTypes.clear();
+        resultTypes.clear();
         emittingReturn = false;
         writer.outdent().write(")").eol();
     }
@@ -139,4 +156,8 @@ public class DisassemblyTypeSectionListener extends BaseDisassemblyListener impl
         }
         writer.write(")").eol();
     }
+
+    public WasmHollowFunctionType[] getFunctionTypes() {
+        return functionTypes.toArray(new WasmHollowFunctionType[0]);
+    }
 }

core/src/main/java/org/teavm/backend/wasm/parser/CodeParser.java

@@ -31,16 +31,11 @@ import org.teavm.backend.wasm.model.expression.WasmSignedType;
 public class CodeParser extends BaseSectionParser {
     private CodeListener codeListener;
     private final List<Block> blockStack = new ArrayList<>();
-    private int functionIndexOffset;
 
     public void setCodeListener(CodeListener codeListener) {
         this.codeListener = codeListener;
     }
 
-    public void setFunctionIndexOffset(int functionIndexOffset) {
-        this.functionIndexOffset = functionIndexOffset;
-    }
-
     @Override
     protected void parseContent() {
         parseCode();
@@ -104,7 +99,7 @@ public class CodeParser extends BaseSectionParser {
                 codeListener.opcode(Opcode.RETURN);
                 break;
             case 0x10:
-                codeListener.call(readLEB() + functionIndexOffset);
+                codeListener.call(readLEB());
                 break;
             case 0x11:
                 codeListener.indirectCall(readLEB(), readLEB());
@@ -598,7 +593,7 @@ public class CodeParser extends BaseSectionParser {
                 break;
 
             case 0xD2:
-                codeListener.functionReference(readLEB() + functionIndexOffset);
+                codeListener.functionReference(readLEB());
                 break;
 
             case 0xD3:

core/src/main/java/org/teavm/backend/wasm/parser/CodeSectionParser.java

@@ -49,7 +49,6 @@ public class CodeSectionParser extends BaseSectionParser {
             var codeListener = listener.code();
             if (codeListener != null) {
                 codeParser.setCodeListener(codeListener);
-                codeParser.setFunctionIndexOffset(functionIndexOffset);
                 codeParser.parse(reader);
             }
         }

core/src/main/java/org/teavm/backend/wasm/parser/FunctionSectionListener.java

@@ -0,0 +1,20 @@
+/*
+ *  Copyright 2024 konsoletyper.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.backend.wasm.parser;
+
+public interface FunctionSectionListener {
+    void function(int index, int typeIndex);
+}

core/src/main/java/org/teavm/backend/wasm/parser/FunctionSectionParser.java

@@ -0,0 +1,34 @@
+/*
+ *  Copyright 2024 konsoletyper.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.backend.wasm.parser;
+
+public class FunctionSectionParser extends BaseSectionParser {
+    private final FunctionSectionListener listener;
+
+    public FunctionSectionParser(FunctionSectionListener listener) {
+        this.listener = listener;
+    }
+
+    @Override
+    protected void parseContent() {
+        var count = reader.readLEB();
+        for (var i = 0; i < count; ++i) {
+            reportAddress();
+            var typeIndex = reader.readLEB();
+            listener.function(i, typeIndex);
+        }
+    }
+}

core/src/main/java/org/teavm/backend/wasm/parser/GlobalSectionParser.java

@@ -18,17 +18,12 @@ package org.teavm.backend.wasm.parser;
 public class GlobalSectionParser extends BaseSectionParser {
     private final GlobalSectionListener listener;
     private CodeParser codeParser;
-    private int functionIndexOffset;
 
     public GlobalSectionParser(GlobalSectionListener listener) {
         this.listener = listener;
         codeParser = new CodeParser();
     }
 
-    public void setFunctionIndexOffset(int functionIndexOffset) {
-        this.functionIndexOffset = functionIndexOffset;
-    }
-
     @Override
     protected void parseContent() {
         var count = readLEB();
@@ -41,7 +36,6 @@ public class GlobalSectionParser extends BaseSectionParser {
                 codeListener = CodeListener.EMPTY;
             }
             codeParser.setCodeListener(codeListener);
-            codeParser.setFunctionIndexOffset(functionIndexOffset);
             if (!codeParser.parseSingleExpression(reader)) {
                 throw new ParseException("Error parsing global initializer", reader.ptr);
             }

core/src/main/java/org/teavm/backend/wasm/parser/WasmHollowFunctionType.java

@@ -0,0 +1,26 @@
+/*
+ *  Copyright 2024 konsoletyper.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.backend.wasm.parser;
+
+public class WasmHollowFunctionType {
+    public final WasmHollowType[] parameterTypes;
+    public final WasmHollowType[] returnTypes;
+
+    public WasmHollowFunctionType(WasmHollowType[] parameterTypes, WasmHollowType[] returnTypes) {
+        this.parameterTypes = parameterTypes;
+        this.returnTypes = returnTypes;
+    }
+}