Multiple memory-retention improvements in IR.

* Lazily allocate Vertex edge sets only when needed for write.
* Use lazily-allocated IntHashMap for IRMethod signatures.
* Clear unused IRScope data strucures after successful JIT.

Also changed some directed graph classes to use Iterable so we
have a free hand to improve the actual data structure used.

Author: headius

core/src/main/java/org/jruby/compiler/JITCompiler.java

@@ -48,6 +48,7 @@
 import org.jruby.util.JavaNameMangler;
 import org.jruby.util.OneShotClassLoader;
 import org.jruby.util.cli.Options;
+import org.jruby.util.collections.IntHashMap;
 import org.jruby.util.log.Logger;
 import org.jruby.util.log.LoggerFactory;
 import org.objectweb.asm.Opcodes;
@@ -58,6 +59,7 @@
 import java.lang.reflect.Method;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.ExecutorService;
@@ -243,7 +245,10 @@ public void run() {
                     log(method, className + "." + methodName, "done jitting");
                 }
 
-                Map<Integer, MethodType> signatures = ((IRMethod)method.getIRMethod()).getNativeSignatures();
+                // clear out data structures we won't need anymore
+                method.getIRMethod().clearAfterJIT();
+
+                IntHashMap<MethodType> signatures = ((IRMethod)method.getIRMethod()).getNativeSignatures();
                 String jittedName = ((IRMethod)method.getIRMethod()).getJittedName();
                 if (signatures.size() == 1) {
                     // only variable-arity
@@ -255,7 +260,7 @@ public void run() {
                                     method.getImplementationClass()));
                 } else {
                     // also specific-arity
-                    for (Map.Entry<Integer, MethodType> entry : signatures.entrySet()) {
+                    for (IntHashMap.Entry<MethodType> entry : signatures.entrySet()) {
                         if (entry.getKey() == -1) continue; // variable arity handle pushed above
 
                         method.switchToJitted(

core/src/main/java/org/jruby/ir/IRMethod.java

@@ -1,5 +1,6 @@
 package org.jruby.ir;
 
+import org.jcodings.util.IntHash;
 import org.jruby.internal.runtime.methods.IRMethodArgs;
 import org.jruby.ir.instructions.Instr;
 import org.jruby.ir.instructions.ReceiveArgBase;
@@ -13,6 +14,7 @@
 import org.jruby.ir.operands.Splat;
 import org.jruby.util.KeyValuePair;
 import org.jruby.parser.StaticScope;
+import org.jruby.util.collections.IntHashMap;
 
 import java.lang.invoke.MethodType;
 import java.util.ArrayList;
@@ -35,7 +37,10 @@ public class IRMethod extends IRScope {
     private List<String[]> argDesc;
 
     // Signatures to the jitted versions of this method
-    private Map<Integer, MethodType> signatures;
+    private IntHashMap<MethodType> signatures = NULL_SIGNATURES;
+
+    // Max arity we support; this is only an artificial number to keep the signature store compact.
+    private static final IntHashMap<MethodType> NULL_SIGNATURES = new IntHashMap<>(0);
 
     // Method name in the jitted version of this method
     private String jittedName;
@@ -48,7 +53,6 @@ public IRMethod(IRManager manager, IRScope lexicalParent, String name,
         this.callArgs = new ArrayList<>();
         this.keywordArgs = new ArrayList<>();
         this.argDesc = new ArrayList<>();
-        this.signatures = new HashMap<>();
 
         if (!getManager().isDryRun() && staticScope != null) {
             staticScope.setIRScope(this);
@@ -115,15 +119,16 @@ public LocalVariable getLocalVariable(String name, int scopeDepth) {
     }
 
     public void addNativeSignature(int arity, MethodType signature) {
+        ensureSignatures();
         signatures.put(arity, signature);
     }
 
-    public MethodType getNativeSignature(int arity) {
-        return signatures.get(arity);
+    public IntHashMap<MethodType> getNativeSignatures() {
+        return signatures;
     }
 
-    public Map<Integer, MethodType> getNativeSignatures() {
-        return Collections.unmodifiableMap(signatures);
+    private void ensureSignatures() {
+        if (signatures == NULL_SIGNATURES) signatures = new IntHashMap<>(5, 0.25f);
     }
 
     public String getJittedName() {

core/src/main/java/org/jruby/ir/IRScope.java

@@ -1208,4 +1208,22 @@ public void savePersistenceInfo(int offset, IRReaderDecoder file) {
         instructionsOffsetInfoPersistenceBuffer = offset;
         persistenceStore = file;
     }
+
+    /**
+     * Flush out data structures no longer neecessary after a successful JIT.
+     */
+    public void clearAfterJIT() {
+        // Clear everything out, since we're never coming back this way again.
+        dfProbs = null;
+        cfg = null;
+        linearizedBBList = null;
+//        localVars = null;
+        definedLocalVars = null;
+        evalScopeVars = null;
+        executedPasses = null;
+//        nextVarIndex = null;
+        instrList = null;
+        nestedClosures = null;
+        lexicalChildren = null;
+    }
 }

core/src/main/java/org/jruby/ir/targets/JVMVisitor.java

@@ -29,6 +29,7 @@
 import org.jruby.util.KeyValuePair;
 import org.jruby.util.RegexpOptions;
 import org.jruby.util.cli.Options;
+import org.jruby.util.collections.IntHashMap;
 import org.jruby.util.log.Logger;
 import org.jruby.util.log.LoggerFactory;
 import org.objectweb.asm.Handle;
@@ -922,7 +923,7 @@ public void DefineClassMethodInstr(DefineClassMethodInstr defineclassmethodinstr
 
         emitMethod(method);
 
-        Map<Integer, MethodType> signatures = method.getNativeSignatures();
+        IntHashMap<MethodType> signatures = method.getNativeSignatures();
 
         MethodType signature = signatures.get(-1);
 
@@ -951,7 +952,7 @@ public void DefineInstanceMethodInstr(DefineInstanceMethodInstr defineinstanceme
         m.loadContext();
 
         emitMethod(method);
-        Map<Integer, MethodType> signatures = method.getNativeSignatures();
+        IntHashMap<MethodType> signatures = method.getNativeSignatures();
 
         MethodType variable = signatures.get(-1); // always a variable arity handle
 
@@ -970,7 +971,7 @@ public void DefineInstanceMethodInstr(DefineInstanceMethodInstr defineinstanceme
         a.invokestatic(p(IRRuntimeHelpers.class), "defCompiledInstanceMethod", defSignature);
     }
 
-    public String pushHandlesForDef(String name, Map<Integer, MethodType> signatures, MethodType variable, String variableOnly, String variableAndSpecific) {
+    public String pushHandlesForDef(String name, IntHashMap<MethodType> signatures, MethodType variable, String variableOnly, String variableAndSpecific) {
         String defSignature;
 
         jvmMethod().pushHandle(new Handle(Opcodes.H_INVOKESTATIC, jvm.clsData().clsName, name, sig(variable.returnType(), variable.parameterArray())));
@@ -981,7 +982,7 @@ public String pushHandlesForDef(String name, Map<Integer, MethodType> signatures
             defSignature = variableAndSpecific;
 
             // FIXME: only supports one arity
-            for (Map.Entry<Integer, MethodType> entry : signatures.entrySet()) {
+            for (IntHashMap.Entry<MethodType> entry : signatures.entrySet()) {
                 if (entry.getKey() == -1) continue; // variable arity signature pushed above
                 jvmMethod().pushHandle(new Handle(Opcodes.H_INVOKESTATIC, jvm.clsData().clsName, name, sig(entry.getValue().returnType(), entry.getValue().parameterArray())));
                 jvmAdapter().pushInt(entry.getKey());

core/src/main/java/org/jruby/ir/util/DataIterable.java

@@ -12,12 +12,12 @@
  * @author enebo
  */
 public class DataIterable<T> implements Iterable<T> {
-    private Set<Edge<T>> edges;
+    private Iterable<Edge<T>> edges;
     private Object type;
     private boolean negate;
     private boolean source;
 
-    public DataIterable(Set<Edge<T>> edges, Object type, boolean source, boolean negate) {
+    public DataIterable(Iterable<Edge<T>> edges, Object type, boolean source, boolean negate) {
         this.edges = edges;
         this.type = type;
         this.negate = negate;

core/src/main/java/org/jruby/ir/util/DataIterator.java

@@ -18,7 +18,7 @@ public class DataIterator<T> implements Iterator<T> {
     private boolean source;
     private boolean negate;
 
-    public DataIterator(Set<Edge<T>> edges, Object type, boolean source, boolean negate) {
+    public DataIterator(Iterable<Edge<T>> edges, Object type, boolean source, boolean negate) {
         this.internalIterator = edges.iterator();
         this.type = type;
         this.source = source;

core/src/main/java/org/jruby/ir/util/EdgeTypeIterable.java

@@ -1,21 +1,22 @@
 package org.jruby.ir.util;
 
 import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 
 /**
  */
 public class EdgeTypeIterable<T> implements Iterable<Edge<T>> {
-    private Set<Edge<T>> edges;
+    private Iterable<Edge<T>> edges;
     private Object type;
     private boolean negate;
 
-    public EdgeTypeIterable(Set<Edge<T>> edges, Object type) {
+    public EdgeTypeIterable(Iterable<Edge<T>> edges, Object type) {
         this(edges, type, false);
 
     }
 
-    public EdgeTypeIterable(Set<Edge<T>> edges, Object type, boolean negate) {
+    public EdgeTypeIterable(Iterable<Edge<T>> edges, Object type, boolean negate) {
         this.edges = edges;
         this.type = type;
         this.negate = negate;

core/src/main/java/org/jruby/ir/util/EdgeTypeIterator.java

@@ -1,6 +1,7 @@
 package org.jruby.ir.util;
 
 import java.util.Iterator;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Set;
 
@@ -13,7 +14,7 @@ public class EdgeTypeIterator<T> implements Iterator<Edge<T>> {
     private Edge nextEdge = null;
     private boolean negate;
 
-    public EdgeTypeIterator(Set<Edge<T>> edges, Object type, boolean negate) {
+    public EdgeTypeIterator(Iterable<Edge<T>> edges, Object type, boolean negate) {
         this.internalIterator = edges.iterator();
         this.type = type;
         this.negate = negate;

core/src/main/java/org/jruby/ir/util/Vertex.java

@@ -8,8 +8,8 @@
 public class Vertex<T> implements Comparable<Vertex<T>> {
     private DirectedGraph graph;
     private T data;
-    private Set<Edge<T>> incoming = null;
-    private Set<Edge<T>> outgoing = null;
+    private Set<Edge<T>> incoming = Collections.EMPTY_SET;
+    private Set<Edge<T>> outgoing = Collections.EMPTY_SET;
     int id;
 
     public Vertex(DirectedGraph graph, T data, int id) {
@@ -24,8 +24,8 @@ public void addEdgeTo(Vertex destination) {
 
     public void addEdgeTo(Vertex destination, Object type) {
         Edge edge = new Edge<T>(this, destination, type);
-        getOutgoingEdges().add(edge);
-        destination.getIncomingEdges().add(edge);
+        addOutgoingEdge(edge);
+        destination.addIncomingEdge(edge);
         graph.edges().add(edge);
     }
 
@@ -42,14 +42,14 @@ public void addEdgeTo(T destination, Object type) {
     public boolean removeEdgeTo(Vertex destination) {
         for (Edge edge: getOutgoingEdges()) {
             if (edge.getDestination() == destination) {
-                getOutgoingEdges().remove(edge);
-                edge.getDestination().getIncomingEdges().remove(edge);
+                removeOutgoingEdge(edge);
+                edge.getDestination().removeIncomingEdge(edge);
                 graph.edges().remove(edge);
                 if(outDegree() == 0) {
-                    outgoing = null;
+                    outgoing = Collections.EMPTY_SET;
                 }
                 if(destination.inDegree() == 0) {
-                    destination.incoming = null;
+                    destination.incoming = Collections.EMPTY_SET;
                 }
                 return true;
             }
@@ -60,18 +60,18 @@ public boolean removeEdgeTo(Vertex destination) {
 
     public void removeAllIncomingEdges() {
         for (Edge edge: getIncomingEdges()) {
-            edge.getSource().getOutgoingEdges().remove(edge);
+            edge.getSource().removeOutgoingEdge(edge);
             graph.edges().remove(edge);
         }
-        incoming = null;
+        incoming = Collections.EMPTY_SET;
     }
 
     public void removeAllOutgoingEdges() {
         for (Edge edge: getOutgoingEdges()) {
-            edge.getDestination().getIncomingEdges().remove(edge);
+            edge.getDestination().removeIncomingEdge(edge);
             graph.edges().remove(edge);
         }
-        outgoing = null;
+        outgoing = Collections.EMPTY_SET;
     }
 
     public void removeAllEdges() {
@@ -80,11 +80,11 @@ public void removeAllEdges() {
     }
 
     public int inDegree() {
-        return (incoming == null) ? 0 : incoming.size();
+        return incoming.size();
     }
 
     public int outDegree() {
-        return (outgoing == null) ? 0 : outgoing.size();
+        return outgoing.size();
     }
 
     public Iterable<Edge<T>> getIncomingEdgesOfType(Object type) {
@@ -172,13 +172,41 @@ public Edge<T> getOutgoingEdge() {
     }
 
     public Set<Edge<T>> getIncomingEdges() {
-        if (incoming == null) incoming = new HashSet<Edge<T>>();
+        return incoming;
+    }
+
+    public void addIncomingEdge(Edge<T> edge) {
+        getIncomingEdgesForWrite().add(edge);
+    }
+
+    public void removeIncomingEdge(Edge<T> edge) {
+        if (incoming.size() == 0) return;
+        incoming.remove(edge);
+        if (incoming.size() == 0) incoming = Collections.EMPTY_SET;
+    }
+
+    private Set<Edge<T>> getIncomingEdgesForWrite() {
+        if (incoming == Collections.EMPTY_SET) incoming = new HashSet<Edge<T>>();
 
         return incoming;
     }
 
+    public void addOutgoingEdge(Edge<T> edge) {
+        getOutgoingEdgesForWrite().add(edge);
+    }
+
+    public void removeOutgoingEdge(Edge<T> edge) {
+        if (outgoing.size() == 0) return;
+        outgoing.remove(edge);
+        if (outgoing.size() == 0) outgoing = Collections.EMPTY_SET;
+    }
+
     public Set<Edge<T>> getOutgoingEdges() {
-        if (outgoing == null) outgoing = new HashSet<Edge<T>>();
+        return outgoing;
+    }
+
+    private Set<Edge<T>> getOutgoingEdgesForWrite() {
+        if (outgoing == Collections.EMPTY_SET) outgoing = new HashSet<Edge<T>>();
 
         return outgoing;
     }