From a9d46ac55ebf1baade61dfd6f8f76dfcd8b9fa4a Mon Sep 17 00:00:00 2001
From: Alexey Andreev <konsoletyper@gmail.com>
Date: Wed, 4 Sep 2024 14:58:56 +0200
Subject: [PATCH] wasm gc: fix issue with type inference for null values

---
 .../gc/methods/WasmGCMethodGenerator.java     |  7 ++++++
 .../model/analysis/BaseTypeInference.java     | 25 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/core/src/main/java/org/teavm/backend/wasm/generate/gc/methods/WasmGCMethodGenerator.java b/core/src/main/java/org/teavm/backend/wasm/generate/gc/methods/WasmGCMethodGenerator.java
index fa806ea16..9816f89f9 100644
--- a/core/src/main/java/org/teavm/backend/wasm/generate/gc/methods/WasmGCMethodGenerator.java
+++ b/core/src/main/java/org/teavm/backend/wasm/generate/gc/methods/WasmGCMethodGenerator.java
@@ -243,6 +243,13 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
                 }
             }
         }
+        for (var i = 0; i < method.getProgram().variableCount(); ++i) {
+            var variable = method.getProgram().variableAt(i);
+            var varNodeIndex = variable.getRegister() >= 0 ? originalIndexToIndex[variable.getRegister()] : -1;
+            if (varNodeIndex >= 0 && variableRepresentatives[varNodeIndex] < 0) {
+                variableRepresentatives[varNodeIndex] = variable.getIndex();
+            }
+        }
 
         for (var i = firstVar; i < ast.getVariables().size(); ++i) {
             var localVar = ast.getVariables().get(i);
diff --git a/core/src/main/java/org/teavm/model/analysis/BaseTypeInference.java b/core/src/main/java/org/teavm/model/analysis/BaseTypeInference.java
index 62ce3840a..538219e40 100644
--- a/core/src/main/java/org/teavm/model/analysis/BaseTypeInference.java
+++ b/core/src/main/java/org/teavm/model/analysis/BaseTypeInference.java
@@ -54,6 +54,7 @@ import org.teavm.model.instructions.NegateInstruction;
 import org.teavm.model.instructions.NullCheckInstruction;
 import org.teavm.model.instructions.NullConstantInstruction;
 import org.teavm.model.instructions.NumericOperandType;
+import org.teavm.model.instructions.PutElementInstruction;
 import org.teavm.model.instructions.PutFieldInstruction;
 import org.teavm.model.instructions.StringConstantInstruction;
 import org.teavm.model.instructions.UnwrapArrayInstruction;
@@ -64,6 +65,7 @@ public abstract class BaseTypeInference<T> {
     private Object[] types;
     private Graph graph;
     private Graph arrayGraph;
+    private Graph backArrayGraph;
     private Graph arrayUnwrapGraph;
     private boolean phisSkipped;
     private boolean backPropagation;
@@ -112,6 +114,7 @@ public abstract class BaseTypeInference<T> {
         }
         graph = visitor.graphBuilder.build();
         arrayGraph = visitor.arrayGraphBuilder.build();
+        backArrayGraph = visitor.backArrayGraphBuilder.build();
         arrayUnwrapGraph = visitor.arrayUnwrapGraphBuilder.build();
     }
 
@@ -191,6 +194,12 @@ public abstract class BaseTypeInference<T> {
                         stack.push(i);
                     }
                 }
+                for (var j : backArrayGraph.incomingEdges(i)) {
+                    if (!nullTypes[j]) {
+                        typeStack.push(elementType((T) types[j]));
+                        stack.push(i);
+                    }
+                }
             }
         }
 
@@ -228,6 +237,15 @@ public abstract class BaseTypeInference<T> {
                     }
                 }
             }
+            if (backArrayGraph.outgoingEdgesCount(variable) > 0) {
+                var elementType = elementType(type);
+                for (var succ : backArrayGraph.outgoingEdges(variable)) {
+                    if (!Objects.equals(types[succ], elementType)) {
+                        stack.push(succ);
+                        typeStack.push(elementType);
+                    }
+                }
+            }
         }
     }
 
@@ -282,11 +300,13 @@ public abstract class BaseTypeInference<T> {
     private class InitialTypeVisitor extends AbstractInstructionVisitor {
         private GraphBuilder graphBuilder;
         private GraphBuilder arrayGraphBuilder;
+        private GraphBuilder backArrayGraphBuilder;
         private GraphBuilder arrayUnwrapGraphBuilder;
 
         InitialTypeVisitor(int size) {
             graphBuilder = new GraphBuilder(size);
             arrayGraphBuilder = new GraphBuilder(size);
+            backArrayGraphBuilder = new GraphBuilder(size);
             arrayUnwrapGraphBuilder = new GraphBuilder(size);
         }
 
@@ -428,6 +448,11 @@ public abstract class BaseTypeInference<T> {
             arrayGraphBuilder.addEdge(insn.getArray().getIndex(), insn.getReceiver().getIndex());
         }
 
+        @Override
+        public void visit(PutElementInstruction insn) {
+            backArrayGraphBuilder.addEdge(insn.getArray().getIndex(), insn.getValue().getIndex());
+        }
+
         @Override
         public void visit(AssignInstruction insn) {
             graphBuilder.addEdge(insn.getAssignee().getIndex(), insn.getReceiver().getIndex());