From 717bbf4a57dc6e8dc30315491691d286894d77a3 Mon Sep 17 00:00:00 2001
From: Alexey Andreev <konsoletyper@gmail.com>
Date: Sun, 5 Nov 2023 22:41:57 +0100
Subject: [PATCH] JS: strip unused functions from hand-written runtime

---
 .../backend/javascript/JavaScriptTarget.java  |  12 +-
 .../javascript/rendering/RuntimeRenderer.java |  56 ++++++-
 .../javascript/templating/AstRemoval.java     |  52 +++++++
 .../templating/RemovablePartsFinder.java      | 144 ++++++++++++++++++
 4 files changed, 252 insertions(+), 12 deletions(-)
 create mode 100644 core/src/main/java/org/teavm/backend/javascript/templating/AstRemoval.java
 create mode 100644 core/src/main/java/org/teavm/backend/javascript/templating/RemovablePartsFinder.java

diff --git a/core/src/main/java/org/teavm/backend/javascript/JavaScriptTarget.java b/core/src/main/java/org/teavm/backend/javascript/JavaScriptTarget.java
index e78701332..337e310e1 100644
--- a/core/src/main/java/org/teavm/backend/javascript/JavaScriptTarget.java
+++ b/core/src/main/java/org/teavm/backend/javascript/JavaScriptTarget.java
@@ -462,15 +462,13 @@ public class JavaScriptTarget implements TeaVMTarget, TeaVMJavaScriptHost {
         int start = sourceWriter.getOffset();
 
         RuntimeRenderer runtimeRenderer = new RuntimeRenderer(classes, sourceWriter);
+        runtimeRenderer.prepareAstParts(renderer.isThreadLibraryUsed());
+        declarations.replay(runtimeRenderer.sink, RememberedSource.FILTER_REF);
+        epilogue.replay(runtimeRenderer.sink, RememberedSource.FILTER_REF);
+        runtimeRenderer.removeUnusedParts();
         runtimeRenderer.renderRuntime();
-        runtimeRenderer.renderHandWrittenRuntime("long.js");
-        if (renderer.isThreadLibraryUsed()) {
-            runtimeRenderer.renderHandWrittenRuntime("thread.js");
-        } else {
-            runtimeRenderer.renderHandWrittenRuntime("simpleThread.js");
-        }
         declarations.write(sourceWriter, 0);
-        runtimeRenderer.renderHandWrittenRuntime("array.js");
+        runtimeRenderer.renderEpilogue();
         epilogue.write(sourceWriter, 0);
 
         printWrapperEnd(sourceWriter);
diff --git a/core/src/main/java/org/teavm/backend/javascript/rendering/RuntimeRenderer.java b/core/src/main/java/org/teavm/backend/javascript/rendering/RuntimeRenderer.java
index 16fd1359f..771eec49d 100644
--- a/core/src/main/java/org/teavm/backend/javascript/rendering/RuntimeRenderer.java
+++ b/core/src/main/java/org/teavm/backend/javascript/rendering/RuntimeRenderer.java
@@ -20,16 +20,24 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.Reader;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
 import org.mozilla.javascript.CompilerEnvirons;
 import org.mozilla.javascript.Context;
 import org.mozilla.javascript.ast.AstRoot;
 import org.teavm.backend.javascript.codegen.SourceWriter;
+import org.teavm.backend.javascript.codegen.SourceWriterSink;
+import org.teavm.backend.javascript.templating.AstRemoval;
+import org.teavm.backend.javascript.templating.RemovablePartsFinder;
 import org.teavm.backend.javascript.templating.TemplatingAstTransformer;
 import org.teavm.backend.javascript.templating.TemplatingAstWriter;
 import org.teavm.model.ClassReaderSource;
 import org.teavm.vm.RenderingException;
 
 public class RuntimeRenderer {
+    private final List<AstRoot> runtimeAstParts = new ArrayList<>();
+    private final List<AstRoot> epilogueAstParts = new ArrayList<>();
+    private final RemovablePartsFinder removablePartsFinder = new RemovablePartsFinder();
     private final ClassReaderSource classSource;
     private final SourceWriter writer;
 
@@ -38,15 +46,35 @@ public class RuntimeRenderer {
         this.writer = writer;
     }
 
-    public void renderRuntime() throws RenderingException {
-        renderHandWrittenRuntime("runtime.js");
-        renderHandWrittenRuntime("intern.js");
+    public void prepareAstParts(boolean threadLibraryUsed) {
+        runtimeAstParts.add(prepareAstPart("runtime.js"));
+        runtimeAstParts.add(prepareAstPart("intern.js"));
+        runtimeAstParts.add(prepareAstPart("long.js"));
+        runtimeAstParts.add(prepareAstPart(threadLibraryUsed ? "thread.js" : "simpleThread.js"));
+        epilogueAstParts.add(prepareAstPart("array.js"));
     }
 
-    public void renderHandWrittenRuntime(String name)  {
-        AstRoot ast = parseRuntime(name);
+    public void renderRuntime() {
+        for (var ast : runtimeAstParts) {
+            renderHandWrittenRuntime(ast);
+        }
+    }
+
+    public void renderEpilogue() {
+        for (var ast : epilogueAstParts) {
+            renderHandWrittenRuntime(ast);
+        }
+    }
+
+    private AstRoot prepareAstPart(String name) {
+        var ast = parseRuntime(name);
         ast.visit(new StringConstantElimination());
         new TemplatingAstTransformer(classSource).visit(ast);
+        removablePartsFinder.visit(ast);
+        return ast;
+    }
+
+    private void renderHandWrittenRuntime(AstRoot ast)  {
         var astWriter = new TemplatingAstWriter(writer, null, null);
         astWriter.hoist(ast);
         astWriter.print(ast);
@@ -66,4 +94,22 @@ public class RuntimeRenderer {
             throw new RenderingException(e);
         }
     }
+
+    public final SourceWriterSink sink = new SourceWriterSink() {
+        @Override
+        public SourceWriterSink appendFunction(String name) {
+            removablePartsFinder.markUsedDeclaration(name);
+            return this;
+        }
+    };
+
+    public void removeUnusedParts() {
+        var removal = new AstRemoval(removablePartsFinder.getAllRemovableParts());
+        for (var part : runtimeAstParts) {
+            removal.visit(part);
+        }
+        for (var part : epilogueAstParts) {
+            removal.visit(part);
+        }
+    }
 }
diff --git a/core/src/main/java/org/teavm/backend/javascript/templating/AstRemoval.java b/core/src/main/java/org/teavm/backend/javascript/templating/AstRemoval.java
new file mode 100644
index 000000000..bb1b055e6
--- /dev/null
+++ b/core/src/main/java/org/teavm/backend/javascript/templating/AstRemoval.java
@@ -0,0 +1,52 @@
+/*
+ *  Copyright 2023 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.javascript.templating;
+
+import java.util.Set;
+import org.mozilla.javascript.ast.AstNode;
+import org.mozilla.javascript.ast.ExpressionStatement;
+import org.mozilla.javascript.ast.VariableDeclaration;
+import org.teavm.backend.javascript.ast.AstVisitor;
+
+public class AstRemoval extends AstVisitor {
+    private Set<AstNode> nodes;
+
+    public AstRemoval(Set<AstNode> nodes) {
+        this.nodes = nodes;
+    }
+
+    @Override
+    public void visit(ExpressionStatement node) {
+        if (nodes.contains(node.getExpression())) {
+            replaceWith(null);
+        } else {
+            super.visit(node);
+        }
+    }
+
+    @Override
+    public void visit(VariableDeclaration node) {
+        for (var iter = node.getVariables().iterator(); iter.hasNext();) {
+            var initializer = iter.next();
+            if (nodes.contains(initializer)) {
+                iter.remove();
+            }
+        }
+        if (node.getVariables().isEmpty()) {
+            replaceWith(null);
+        }
+    }
+}
diff --git a/core/src/main/java/org/teavm/backend/javascript/templating/RemovablePartsFinder.java b/core/src/main/java/org/teavm/backend/javascript/templating/RemovablePartsFinder.java
new file mode 100644
index 000000000..9b442f8ba
--- /dev/null
+++ b/core/src/main/java/org/teavm/backend/javascript/templating/RemovablePartsFinder.java
@@ -0,0 +1,144 @@
+/*
+ *  Copyright 2023 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.javascript.templating;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.mozilla.javascript.ast.Assignment;
+import org.mozilla.javascript.ast.AstNode;
+import org.mozilla.javascript.ast.ElementGet;
+import org.mozilla.javascript.ast.ExpressionStatement;
+import org.mozilla.javascript.ast.FunctionNode;
+import org.mozilla.javascript.ast.Name;
+import org.mozilla.javascript.ast.PropertyGet;
+import org.mozilla.javascript.ast.Scope;
+import org.mozilla.javascript.ast.VariableDeclaration;
+import org.teavm.backend.javascript.ast.AstVisitor;
+
+public class RemovablePartsFinder extends AstVisitor {
+    private Map<String, List<AstNode>> removableDeclarations = new HashMap<>();
+    private Map<String, Scope> removableDeclarationScopes = new HashMap<>();
+    private Map<String, Set<String>> dependencies = new HashMap<>();
+    private String insideDeclaration;
+    private boolean topLevel = true;
+
+    @Override
+    public void visit(FunctionNode node) {
+        if (topLevel) {
+            if (node.getName() != null && !node.getName().isEmpty()) {
+                removableDeclarations.computeIfAbsent(node.getName(), k -> new ArrayList<>()).add(node);
+                removableDeclarationScopes.put(node.getName(), scopeOfId(node.getName()));
+            }
+            topLevel = false;
+            insideDeclaration = node.getName();
+            visit(node.getBody());
+            insideDeclaration = null;
+            topLevel = true;
+        } else {
+            super.visit(node);
+        }
+    }
+
+    @Override
+    public void visit(VariableDeclaration node) {
+        if (topLevel) {
+            for (var initializer : node.getVariables()) {
+                var name = extractName(initializer.getTarget());
+                if (name != null) {
+                    removableDeclarations.computeIfAbsent(name.getIdentifier(), k -> new ArrayList<>())
+                            .add(initializer);
+                    removableDeclarationScopes.put(name.getIdentifier(), scopeOfId(name.getIdentifier()));
+                    if (initializer.getInitializer() != null) {
+                        topLevel = false;
+                        insideDeclaration = name.getIdentifier();
+                        visit(initializer.getInitializer());
+                        insideDeclaration = null;
+                        topLevel = true;
+                    }
+                }
+            }
+        } else {
+            super.visit(node);
+        }
+    }
+
+    @Override
+    public void visit(ExpressionStatement node) {
+        if (topLevel && node.getExpression() instanceof Assignment) {
+            var assign = (Assignment) node.getExpression();
+            var name = extractName(assign.getLeft());
+            removableDeclarations.computeIfAbsent(name.getIdentifier(), k -> new ArrayList<>())
+                    .add(node.getExpression());
+            removableDeclarationScopes.put(name.getIdentifier(), scopeOfId(name.getIdentifier()));
+            if (name != null) {
+                topLevel = false;
+                insideDeclaration = name.getIdentifier();
+                visit(assign.getRight());
+                insideDeclaration = null;
+                topLevel = true;
+                return;
+            }
+        }
+        super.visit(node);
+    }
+
+    @Override
+    public void visit(PropertyGet node) {
+        visit(node.getTarget());
+    }
+
+    @Override
+    public void visit(Name node) {
+        if (scopeOfId(node.getIdentifier()) == removableDeclarationScopes.get(node.getIdentifier())
+                && insideDeclaration != null) {
+            dependencies.computeIfAbsent(insideDeclaration, k -> new HashSet<>()).add(node.getIdentifier());
+        }
+    }
+
+    private Name extractName(AstNode node) {
+        if (node instanceof Name) {
+            return (Name) node;
+        } else if (node instanceof PropertyGet) {
+            return extractName(((PropertyGet) node).getTarget());
+        } else if (node instanceof ElementGet) {
+            return extractName(((ElementGet) node).getTarget());
+        } else {
+            return null;
+        }
+    }
+
+    public void markUsedDeclaration(String name) {
+        removableDeclarations.remove(name);
+        var dependenciesToFollow = dependencies.remove(name);
+        if (dependenciesToFollow != null) {
+            for (var dependency : dependenciesToFollow) {
+                markUsedDeclaration(dependency);
+            }
+        }
+    }
+
+    public Set<AstNode> getAllRemovableParts() {
+        var nodes = new HashSet<AstNode>();
+        for (var parts : removableDeclarations.values()) {
+            nodes.addAll(parts);
+        }
+        return nodes;
+    }
+}