Fix broken call graph for nested lambdas

pom.xml

@@ -17,7 +17,7 @@
     <dependency>
       <groupId>org.apache.bcel</groupId>
       <artifactId>bcel</artifactId>
-      <version>6.0</version>
+      <version>6.2</version>
       <scope>provided</scope>
     </dependency>
     <dependency>
@@ -26,6 +26,24 @@
       <version>3.12.1.GA</version>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+        <groupId>io.cucumber</groupId>
+        <artifactId>cucumber-java</artifactId>
+        <version>2.3.1</version>
+        <scope>test</scope>
+    </dependency>
+    <dependency>
+        <groupId>io.cucumber</groupId>
+        <artifactId>cucumber-junit</artifactId>
+        <version>2.3.1</version>
+        <scope>test</scope>
+    </dependency>
+    <dependency>
+        <groupId>junit</groupId>
+        <artifactId>junit</artifactId>
+        <version>4.12</version>
+        <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <plugins>

src/main/java/gr/gousiosg/javacg/stat/ClassVisitor.java

@@ -45,6 +45,7 @@ public class ClassVisitor extends EmptyVisitor {
     private JavaClass clazz;
     private ConstantPoolGen constants;
     private String classReferenceFormat;
+    private final DynamicCallManager DCManager = new DynamicCallManager();
 
     public ClassVisitor(JavaClass jc) {
         clazz = jc;
@@ -55,8 +56,13 @@ public ClassVisitor(JavaClass jc) {
     public void visitJavaClass(JavaClass jc) {
         jc.getConstantPool().accept(this);
         Method[] methods = jc.getMethods();
-        for (int i = 0; i < methods.length; i++)
-            methods[i].accept(this);
+        for (int i = 0; i < methods.length; i++) {
+            Method method = methods[i];
+            DCManager.retrieveCalls(method, jc);
+            DCManager.linkCalls(method);
+            method.accept(this);
+
+        }
     }
 
     public void visitConstantPool(ConstantPool constantPool) {

src/main/java/gr/gousiosg/javacg/stat/DynamicCallManager.java

@@ -0,0 +1,116 @@
+package gr.gousiosg.javacg.stat;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import org.apache.bcel.classfile.Attribute;
+import org.apache.bcel.classfile.BootstrapMethod;
+import org.apache.bcel.classfile.BootstrapMethods;
+import org.apache.bcel.classfile.ConstantCP;
+import org.apache.bcel.classfile.ConstantMethodHandle;
+import org.apache.bcel.classfile.ConstantNameAndType;
+import org.apache.bcel.classfile.ConstantPool;
+import org.apache.bcel.classfile.ConstantUtf8;
+import org.apache.bcel.classfile.JavaClass;
+import org.apache.bcel.classfile.Method;
+
+/**
+ * {@link DynamicCallManager} provides facilities to retrieve information about
+ * dynamic calls statically.
+ * <p>
+ * Most of the time, call relationships are explicit, which allows to properly
+ * build the call graph statically. But in the case of dynamic linking, i.e.
+ * <code>invokedynamic</code> instructions, this relationship might be unknown
+ * until the code is actually executed. Indeed, bootstrap methods are used to
+ * dynamically link the code at first call. One can read details about the
+ * <a href=
+ * "https://docs.oracle.com/javase/8/docs/technotes/guides/vm/multiple-language-support.html#invokedynamic"><code>invokedynamic</code>
+ * instruction</a> to know more about this mechanism.
+ * <p>
+ * Nested lambdas are particularly subject to such absence of concrete caller,
+ * which lead us to produce method names like <code>lambda$null$0</code>, which
+ * breaks the call graph. This information can however be retrieved statically
+ * through the code of the bootstrap method called.
+ * <p>
+ * In {@link #retrieveCalls(Method, JavaClass)}, we retrieve the (called,
+ * caller) relationships by analyzing the code of the caller {@link Method}.
+ * This information is then used in {@link #linkCalls(Method)} to rename the
+ * called {@link Method} properly.
+ * 
+ * @author Matthieu Vergne <matthieu.vergne@gmail.com>
+ *
+ */
+public class DynamicCallManager {
+    private static final Pattern BOOTSTRAP_CALL_PATTERN = Pattern
+            .compile("invokedynamic\t(\\d+):\\S+ \\S+ \\(\\d+\\)");
+    private static final int CALL_HANDLE_INDEX_ARGUMENT = 1;
+
+    private final Map<String, String> dynamicCallers = new HashMap<>();
+
+    /**
+     * Retrieve dynamic call relationships based on the code of the provided
+     * {@link Method}.
+     * 
+     * @param method
+     *            {@link Method} to analyze the code
+     * @param jc
+     *            {@link JavaClass} info, which contains the bootstrap methods
+     * @see #linkCalls(Method)
+     */
+    public void retrieveCalls(Method method, JavaClass jc) {
+        if (method.isAbstract()) {
+            // No code to consider
+            return;
+        }
+        ConstantPool cp = method.getConstantPool();
+        BootstrapMethod[] boots = getBootstrapMethods(jc);
+        String code = method.getCode().toString();
+        Matcher matcher = BOOTSTRAP_CALL_PATTERN.matcher(code);
+        while (matcher.find()) {
+            int bootIndex = Integer.parseInt(matcher.group(1));
+            BootstrapMethod bootMethod = boots[bootIndex];
+            int calledIndex = bootMethod.getBootstrapArguments()[CALL_HANDLE_INDEX_ARGUMENT];
+            String calledName = getMethodNameFromHandleIndex(cp, calledIndex);
+            String callerName = method.getName();
+            dynamicCallers.put(calledName, callerName);
+        }
+    }
+
+    private String getMethodNameFromHandleIndex(ConstantPool cp, int callIndex) {
+        ConstantMethodHandle handle = (ConstantMethodHandle) cp.getConstant(callIndex);
+        ConstantCP ref = (ConstantCP) cp.getConstant(handle.getReferenceIndex());
+        ConstantNameAndType nameAndType = (ConstantNameAndType) cp.getConstant(ref.getNameAndTypeIndex());
+        return nameAndType.getName(cp);
+    }
+
+    /**
+     * Link the {@link Method}'s name to its concrete caller if required.
+     * 
+     * @param method
+     *            {@link Method} to analyze
+     * @see #retrieveCalls(Method, JavaClass)
+     */
+    public void linkCalls(Method method) {
+        int nameIndex = method.getNameIndex();
+        ConstantPool cp = method.getConstantPool();
+        String methodName = ((ConstantUtf8) cp.getConstant(nameIndex)).getBytes();
+        String linkedName = methodName;
+        String callerName = methodName;
+        while (linkedName.matches("(lambda\\$)+null(\\$\\d+)+")) {
+            callerName = dynamicCallers.get(callerName);
+            linkedName = linkedName.replace("null", callerName);
+        }
+        cp.setConstant(nameIndex, new ConstantUtf8(linkedName));
+    }
+
+    private BootstrapMethod[] getBootstrapMethods(JavaClass jc) {
+        for (Attribute attribute : jc.getAttributes()) {
+            if (attribute instanceof BootstrapMethods) {
+                return ((BootstrapMethods) attribute).getBootstrapMethods();
+            }
+        }
+        return new BootstrapMethod[] {};
+    }
+}

src/test/java/javacg/JARBuilder.java

@@ -0,0 +1,117 @@
+package javacg;
+
+import java.io.BufferedInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.StringWriter;
+import java.net.URI;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.LinkedList;
+import java.util.jar.Attributes;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
+import java.util.jar.Manifest;
+
+import javax.tools.Diagnostic;
+import javax.tools.DiagnosticCollector;
+import javax.tools.JavaCompiler;
+import javax.tools.JavaCompiler.CompilationTask;
+import javax.tools.JavaFileObject;
+import javax.tools.SimpleJavaFileObject;
+import javax.tools.StandardJavaFileManager;
+import javax.tools.StandardLocation;
+import javax.tools.ToolProvider;
+
+public class JARBuilder {
+	private static final String TEMP_DIR = System.getProperty("java.io.tmpdir");
+
+	private final JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
+	private final DiagnosticCollector<JavaFileObject> diagnostics = new DiagnosticCollector<JavaFileObject>();
+	private final StandardJavaFileManager fileManager = compiler.getStandardFileManager(diagnostics, null, null);
+	private final LinkedList<JavaFileObject> compilationUnits = new LinkedList<>();
+	private final Collection<File> classFiles = new LinkedList<>();
+
+	public JARBuilder() throws IOException {
+		fileManager.setLocation(StandardLocation.CLASS_OUTPUT, Arrays.asList(new File(TEMP_DIR)));
+	}
+
+	public void add(String className, String classCode) throws IOException {
+		compilationUnits.add(createJavaFile(className, classCode));
+		classFiles.add(new File(TEMP_DIR, className + ".class"));
+	}
+
+	public File build() throws FileNotFoundException, IOException {
+		CompilationTask task = compiler.getTask(null, fileManager, diagnostics, null, null, compilationUnits);
+		boolean success = task.call();
+		if (!success) {
+			displayDiagnostic(diagnostics);
+			throw new RuntimeException("Cannot compile classes for the JAR");
+		}
+
+		File file = File.createTempFile("test", ".jar");
+		JarOutputStream jar = new JarOutputStream(new FileOutputStream(file), createManifest());
+		for (File classFile : classFiles) {
+			add(classFile, jar);
+		}
+		jar.close();
+		return file;
+	}
+
+	private void displayDiagnostic(DiagnosticCollector<JavaFileObject> diagnostics) {
+		for (Diagnostic<?> diagnostic : diagnostics.getDiagnostics()) {
+			JavaSourceFromString sourceClass = (JavaSourceFromString) diagnostic.getSource();
+			System.err.println("-----");
+			System.err.println("Source: " + sourceClass.getName());
+			System.err.println("Message: " + diagnostic.getMessage(null));
+			System.err.println("Position: " + diagnostic.getPosition());
+			System.err.println(diagnostic.getKind() + " " + diagnostic.getCode());
+		}
+	}
+
+	private JavaFileObject createJavaFile(String className, String classCode) throws IOException {
+		StringWriter writer = new StringWriter();
+		writer.append(classCode);
+		writer.close();
+		JavaFileObject file = new JavaSourceFromString(className, writer.toString());
+		return file;
+	}
+
+	private Manifest createManifest() {
+		Manifest manifest = new Manifest();
+		manifest.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0");
+		return manifest;
+	}
+
+	private void add(File classFile, JarOutputStream jar) throws IOException {
+		JarEntry entry = new JarEntry(classFile.getPath().replace("\\", "/"));
+		jar.putNextEntry(entry);
+		try (BufferedInputStream in = new BufferedInputStream(new FileInputStream(classFile))) {
+			byte[] buffer = new byte[1024];
+			while (true) {
+				int count = in.read(buffer);
+				if (count == -1)
+					break;
+				jar.write(buffer, 0, count);
+			}
+			jar.closeEntry();
+		}
+	}
+}
+
+class JavaSourceFromString extends SimpleJavaFileObject {
+	final String code;
+
+	JavaSourceFromString(String name, String code) throws IOException {
+		super(URI.create("string:///" + name.replace('.', '/') + Kind.SOURCE.extension), Kind.SOURCE);
+		this.code = code;
+	}
+
+	@Override
+	public CharSequence getCharContent(boolean ignoreEncodingErrors) {
+		return code;
+	}
+}

src/test/java/javacg/RunCucumberTest.java

@@ -0,0 +1,10 @@
+package javacg;
+
+import cucumber.api.CucumberOptions;
+import cucumber.api.junit.Cucumber;
+import org.junit.runner.RunWith;
+
+@RunWith(Cucumber.class)
+@CucumberOptions(plugin = { "pretty" })
+public class RunCucumberTest {
+}

src/test/java/javacg/StepDefinitions.java

@@ -0,0 +1,48 @@
+package javacg;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintStream;
+
+import cucumber.api.java.en.Given;
+import cucumber.api.java.en.Then;
+import cucumber.api.java.en.When;
+import gr.gousiosg.javacg.stat.JCallGraph;
+
+public class StepDefinitions {
+	private final JARBuilder jarBuilder;
+	private String result;
+
+	public StepDefinitions() throws IOException {
+		jarBuilder = new JARBuilder();
+	}
+
+	@Given("^I have the class \"([^\"]*)\" with code:$")
+	public void i_have_the_class_with_code(String className, String classCode) throws Exception {
+		jarBuilder.add(className, classCode);
+	}
+
+	@When("^I run the analyze$")
+	public void i_analyze_it() throws Exception {
+		File jarFile = jarBuilder.build();
+
+		PrintStream oldOut = System.out;
+		ByteArrayOutputStream resultBuffer = new ByteArrayOutputStream();
+		System.setOut(new PrintStream(resultBuffer));
+		JCallGraph.main(new String[] { jarFile.getPath() });
+		System.setOut(oldOut);
+
+		result = resultBuffer.toString();
+	}
+
+	@Then("^the result should contain:$")
+	public void the_result_should_contain(String line) throws Exception {
+		if (result.contains(line)) {
+			// OK
+		} else {
+			System.err.println(result);
+			throw new RuntimeException("Cannot found: " + line);
+		}
+	}
+}