From 4f218b17f8d49f37d7fc5db58f527502798768cd Mon Sep 17 00:00:00 2001
From: Alexey Andreev <Alexey.Andreev@jetbrains.com>
Date: Thu, 22 Sep 2016 00:58:21 +0300
Subject: [PATCH] WASM: creating memory access trace transformation

---
 .idea/inspectionProfiles/Project_Default.xml  |   4 +
 .../org/teavm/backend/wasm/WasmTarget.java    |   5 +
 .../model/expression/WasmLoadFloat32.java     |   4 +-
 .../model/expression/WasmLoadFloat64.java     |   4 +-
 .../wasm/model/expression/WasmLoadInt32.java  |   4 +-
 .../wasm/model/expression/WasmLoadInt64.java  |   4 +-
 .../model/expression/WasmMemoryAccess.java    |  22 ++
 .../WasmReplacingExpressionVisitor.java       | 224 ++++++++++++++++++
 .../model/expression/WasmStoreFloat32.java    |   4 +-
 .../model/expression/WasmStoreFloat64.java    |   4 +-
 .../wasm/model/expression/WasmStoreInt32.java |   4 +-
 .../wasm/model/expression/WasmStoreInt64.java |   4 +-
 .../MemoryAccessTraceTransformation.java      |  58 +++++
 samples/hello/teavm-samples-hello.iml         |   3 +
 14 files changed, 340 insertions(+), 8 deletions(-)
 create mode 100644 core/src/main/java/org/teavm/backend/wasm/model/expression/WasmMemoryAccess.java
 create mode 100644 core/src/main/java/org/teavm/backend/wasm/model/expression/WasmReplacingExpressionVisitor.java
 create mode 100644 core/src/main/java/org/teavm/backend/wasm/transformation/MemoryAccessTraceTransformation.java

diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
index be16622b6..abc9bfdca 100644
--- a/.idea/inspectionProfiles/Project_Default.xml
+++ b/.idea/inspectionProfiles/Project_Default.xml
@@ -74,6 +74,10 @@
         <value />
       </option>
     </inspection_tool>
+    <inspection_tool class="MissingOverrideAnnotation" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoreObjectMethods" value="true" />
+      <option name="ignoreAnonymousClassMethods" value="false" />
+    </inspection_tool>
     <inspection_tool class="NonOsgiMavenDependency" enabled="false" level="WARNING" enabled_by_default="false" />
     <inspection_tool class="NullableProblems" enabled="true" level="WEAK WARNING" enabled_by_default="true">
       <option name="REPORT_NULLABLE_METHOD_OVERRIDES_NOTNULL" value="true" />
diff --git a/core/src/main/java/org/teavm/backend/wasm/WasmTarget.java b/core/src/main/java/org/teavm/backend/wasm/WasmTarget.java
index 560c09fb8..5651ee50f 100644
--- a/core/src/main/java/org/teavm/backend/wasm/WasmTarget.java
+++ b/core/src/main/java/org/teavm/backend/wasm/WasmTarget.java
@@ -72,6 +72,7 @@ import org.teavm.backend.wasm.render.WasmBinaryRenderer;
 import org.teavm.backend.wasm.render.WasmBinaryWriter;
 import org.teavm.backend.wasm.render.WasmCRenderer;
 import org.teavm.backend.wasm.render.WasmRenderer;
+import org.teavm.backend.wasm.transformation.MemoryAccessTraceTransformation;
 import org.teavm.dependency.ClassDependency;
 import org.teavm.dependency.DependencyChecker;
 import org.teavm.dependency.DependencyListener;
@@ -341,6 +342,10 @@ public class WasmTarget implements TeaVMTarget {
 
         new UnusedFunctionElimination(module).apply();
 
+        if (Boolean.parseBoolean(System.getProperty("wasm.memoryTrace", "false"))) {
+            new MemoryAccessTraceTransformation(module).apply();
+        }
+
         WasmBinaryWriter writer = new WasmBinaryWriter();
         WasmBinaryRenderer renderer = new WasmBinaryRenderer(writer);
         renderer.render(module);
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat32.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat32.java
index a845ffd5b..93097a25d 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat32.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat32.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmLoadFloat32 extends WasmExpression {
+public class WasmLoadFloat32 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
 
@@ -35,10 +35,12 @@ public class WasmLoadFloat32 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat64.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat64.java
index c594ca36c..7e72cd342 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat64.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadFloat64.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmLoadFloat64 extends WasmExpression {
+public class WasmLoadFloat64 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
 
@@ -35,10 +35,12 @@ public class WasmLoadFloat64 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt32.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt32.java
index ad470429a..6648a25a5 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt32.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt32.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmLoadInt32 extends WasmExpression {
+public class WasmLoadInt32 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
     private WasmInt32Subtype convertFrom;
@@ -38,10 +38,12 @@ public class WasmLoadInt32 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt64.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt64.java
index ba1315c39..addcae63a 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt64.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmLoadInt64.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmLoadInt64 extends WasmExpression {
+public class WasmLoadInt64 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
     private WasmInt64Subtype convertFrom;
@@ -38,10 +38,12 @@ public class WasmLoadInt64 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmMemoryAccess.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmMemoryAccess.java
new file mode 100644
index 000000000..3b7a8b97e
--- /dev/null
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmMemoryAccess.java
@@ -0,0 +1,22 @@
+/*
+ *  Copyright 2016 Alexey Andreev.
+ *
+ *  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.model.expression;
+
+public interface WasmMemoryAccess {
+    WasmExpression getIndex();
+
+    void setIndex(WasmExpression index);
+}
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmReplacingExpressionVisitor.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmReplacingExpressionVisitor.java
new file mode 100644
index 000000000..a665ce10e
--- /dev/null
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmReplacingExpressionVisitor.java
@@ -0,0 +1,224 @@
+/*
+ *  Copyright 2016 Alexey Andreev.
+ *
+ *  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.model.expression;
+
+import java.util.List;
+import java.util.function.Function;
+import org.teavm.backend.wasm.model.WasmFunction;
+
+public class WasmReplacingExpressionVisitor implements WasmExpressionVisitor {
+    private Function<WasmExpression, WasmExpression> mapper;
+
+    public WasmReplacingExpressionVisitor(
+            Function<WasmExpression, WasmExpression> mapper) {
+        this.mapper = mapper;
+    }
+
+    public void replace(WasmFunction function) {
+        replaceExpressions(function.getBody());
+    }
+
+    private void replaceExpressions(List<WasmExpression> expressions) {
+        for (int i = 0; i < expressions.size(); ++i) {
+            WasmExpression part = expressions.get(i);
+            part.acceptVisitor(this);
+            part = mapper.apply(part);
+            expressions.set(i, part);
+        }
+    }
+
+    @Override
+    public void visit(WasmBlock expression) {
+        replaceExpressions(expression.getBody());
+    }
+
+    @Override
+    public void visit(WasmBranch expression) {
+        expression.getCondition().acceptVisitor(this);
+        expression.setCondition(mapper.apply(expression.getCondition()));
+        if (expression.getResult() != null) {
+            expression.getResult().acceptVisitor(this);
+            expression.setResult(mapper.apply(expression.getResult()));
+        }
+    }
+
+    @Override
+    public void visit(WasmBreak expression) {
+        if (expression.getResult() != null) {
+            expression.getResult().acceptVisitor(this);
+            expression.setResult(mapper.apply(expression.getResult()));
+        }
+    }
+
+    @Override
+    public void visit(WasmSwitch expression) {
+        expression.getSelector().acceptVisitor(this);
+        expression.setSelector(mapper.apply(expression.getSelector()));
+    }
+
+    @Override
+    public void visit(WasmConditional expression) {
+        expression.getCondition().acceptVisitor(this);
+        expression.setCondition(mapper.apply(expression.getCondition()));
+        replaceExpressions(expression.getThenBlock().getBody());
+        replaceExpressions(expression.getElseBlock().getBody());
+    }
+
+    @Override
+    public void visit(WasmReturn expression) {
+        if (expression.getValue() != null) {
+            expression.getValue().acceptVisitor(this);
+            expression.setValue(mapper.apply(expression.getValue()));
+        }
+    }
+
+    @Override
+    public void visit(WasmUnreachable expression) {
+    }
+
+    @Override
+    public void visit(WasmInt32Constant expression) {
+    }
+
+    @Override
+    public void visit(WasmInt64Constant expression) {
+    }
+
+    @Override
+    public void visit(WasmFloat32Constant expression) {
+    }
+
+    @Override
+    public void visit(WasmFloat64Constant expression) {
+    }
+
+    @Override
+    public void visit(WasmGetLocal expression) {
+    }
+
+    @Override
+    public void visit(WasmSetLocal expression) {
+        expression.getValue().acceptVisitor(this);
+    }
+
+    @Override
+    public void visit(WasmIntBinary expression) {
+        expression.getFirst().acceptVisitor(this);
+        expression.setFirst(mapper.apply(expression.getFirst()));
+        expression.getSecond().acceptVisitor(this);
+        expression.setSecond(mapper.apply(expression.getSecond()));
+    }
+
+    @Override
+    public void visit(WasmFloatBinary expression) {
+        expression.getFirst().acceptVisitor(this);
+        expression.setFirst(mapper.apply(expression.getFirst()));
+        expression.getSecond().acceptVisitor(this);
+        expression.setSecond(mapper.apply(expression.getSecond()));
+    }
+
+    @Override
+    public void visit(WasmIntUnary expression) {
+        expression.getOperand().acceptVisitor(this);
+        expression.setOperand(mapper.apply(expression.getOperand()));
+    }
+
+    @Override
+    public void visit(WasmFloatUnary expression) {
+        expression.getOperand().acceptVisitor(this);
+        expression.setOperand(mapper.apply(expression.getOperand()));
+    }
+
+    @Override
+    public void visit(WasmConversion expression) {
+        expression.getOperand().acceptVisitor(this);
+        expression.setOperand(mapper.apply(expression.getOperand()));
+    }
+
+    @Override
+    public void visit(WasmCall expression) {
+        replaceExpressions(expression.getArguments());
+    }
+
+    @Override
+    public void visit(WasmIndirectCall expression) {
+        expression.getSelector().acceptVisitor(this);
+        expression.setSelector(mapper.apply(expression.getSelector()));
+        replaceExpressions(expression.getArguments());
+    }
+
+    @Override
+    public void visit(WasmDrop expression) {
+        expression.getOperand().acceptVisitor(this);
+        expression.setOperand(mapper.apply(expression.getOperand()));
+    }
+
+    @Override
+    public void visit(WasmLoadInt32 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+    }
+
+    @Override
+    public void visit(WasmLoadInt64 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+    }
+
+    @Override
+    public void visit(WasmLoadFloat32 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+    }
+
+    @Override
+    public void visit(WasmLoadFloat64 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+    }
+
+    @Override
+    public void visit(WasmStoreInt32 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+        expression.getValue().acceptVisitor(this);
+        expression.setValue(mapper.apply(expression.getValue()));
+    }
+
+    @Override
+    public void visit(WasmStoreInt64 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+        expression.getValue().acceptVisitor(this);
+        expression.setValue(mapper.apply(expression.getValue()));
+    }
+
+    @Override
+    public void visit(WasmStoreFloat32 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+        expression.getValue().acceptVisitor(this);
+        expression.setValue(mapper.apply(expression.getValue()));
+    }
+
+    @Override
+    public void visit(WasmStoreFloat64 expression) {
+        expression.getIndex().acceptVisitor(this);
+        expression.setIndex(mapper.apply(expression.getIndex()));
+        expression.getValue().acceptVisitor(this);
+        expression.setValue(mapper.apply(expression.getValue()));
+    }
+}
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat32.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat32.java
index fb743f48c..d6356813d 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat32.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat32.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmStoreFloat32 extends WasmExpression {
+public class WasmStoreFloat32 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
     private WasmExpression value;
@@ -38,10 +38,12 @@ public class WasmStoreFloat32 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat64.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat64.java
index 812ee9177..758439b93 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat64.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreFloat64.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmStoreFloat64 extends WasmExpression {
+public class WasmStoreFloat64 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
     private WasmExpression value;
@@ -38,10 +38,12 @@ public class WasmStoreFloat64 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt32.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt32.java
index edd6d750f..998f3272a 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt32.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt32.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmStoreInt32 extends WasmExpression {
+public class WasmStoreInt32 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
     private WasmExpression value;
@@ -42,10 +42,12 @@ public class WasmStoreInt32 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt64.java b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt64.java
index f05c2a6da..4156c0c09 100644
--- a/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt64.java
+++ b/core/src/main/java/org/teavm/backend/wasm/model/expression/WasmStoreInt64.java
@@ -17,7 +17,7 @@ package org.teavm.backend.wasm.model.expression;
 
 import java.util.Objects;
 
-public class WasmStoreInt64 extends WasmExpression {
+public class WasmStoreInt64 extends WasmExpression implements WasmMemoryAccess {
     private int alignment;
     private WasmExpression index;
     private WasmExpression value;
@@ -42,10 +42,12 @@ public class WasmStoreInt64 extends WasmExpression {
         this.alignment = alignment;
     }
 
+    @Override
     public WasmExpression getIndex() {
         return index;
     }
 
+    @Override
     public void setIndex(WasmExpression index) {
         Objects.requireNonNull(index);
         this.index = index;
diff --git a/core/src/main/java/org/teavm/backend/wasm/transformation/MemoryAccessTraceTransformation.java b/core/src/main/java/org/teavm/backend/wasm/transformation/MemoryAccessTraceTransformation.java
new file mode 100644
index 000000000..62447854b
--- /dev/null
+++ b/core/src/main/java/org/teavm/backend/wasm/transformation/MemoryAccessTraceTransformation.java
@@ -0,0 +1,58 @@
+/*
+ *  Copyright 2016 Alexey Andreev.
+ *
+ *  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.transformation;
+
+import org.teavm.backend.wasm.model.WasmFunction;
+import org.teavm.backend.wasm.model.WasmModule;
+import org.teavm.backend.wasm.model.WasmType;
+import org.teavm.backend.wasm.model.expression.WasmCall;
+import org.teavm.backend.wasm.model.expression.WasmInt32Constant;
+import org.teavm.backend.wasm.model.expression.WasmMemoryAccess;
+import org.teavm.backend.wasm.model.expression.WasmReplacingExpressionVisitor;
+
+public class MemoryAccessTraceTransformation {
+    private WasmModule module;
+
+    public MemoryAccessTraceTransformation(WasmModule module) {
+        this.module = module;
+    }
+
+    public void apply() {
+        WasmFunction traceFunction = new WasmFunction("traceMemoryAccess");
+        traceFunction.setImportModule("debug");
+        traceFunction.setImportName("traceMemoryAccess");
+        traceFunction.getParameters().add(WasmType.INT32);
+        traceFunction.getParameters().add(WasmType.INT32);
+        traceFunction.setResult(WasmType.INT32);
+        module.add(traceFunction);
+
+        int[] positionHolder = { 0 };
+        WasmReplacingExpressionVisitor visitor = new WasmReplacingExpressionVisitor(expression -> {
+            if (expression instanceof WasmMemoryAccess) {
+                WasmMemoryAccess memoryAccess = (WasmMemoryAccess) expression;
+                WasmCall call = new WasmCall(traceFunction.getName());
+                call.setImported(true);
+                call.getArguments().add(new WasmInt32Constant(positionHolder[0]++));
+                call.getArguments().add(memoryAccess.getIndex());
+                memoryAccess.setIndex(call);
+            }
+            return expression;
+        });
+        for (WasmFunction function : module.getFunctions().values()) {
+            visitor.replace(function);
+        }
+    }
+}
diff --git a/samples/hello/teavm-samples-hello.iml b/samples/hello/teavm-samples-hello.iml
index 3f2c66598..6070666aa 100644
--- a/samples/hello/teavm-samples-hello.iml
+++ b/samples/hello/teavm-samples-hello.iml
@@ -10,6 +10,9 @@
           <root url="file://$MODULE_DIR$/target/generated/js" relative="/" />
           <root url="file://$MODULE_DIR$/src/main/webapp" relative="/" />
         </webroots>
+        <sourceRoots>
+          <root url="file://$MODULE_DIR$/src/main/java" />
+        </sourceRoots>
       </configuration>
     </facet>
   </component>