From 0938b3ccfaee07e26b2ff83c5dafb6829d171bdb Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Tue, 23 Apr 2019 15:59:42 +0800
Subject: [PATCH] mc support JPMS; merge JavaFileObjects from specified
 ClassLoader. #661

---
 .../compiler/DynamicJavaFileManager.java      | 72 +++++++++++++++----
 1 file changed, 57 insertions(+), 15 deletions(-)

diff --git a/memorycompiler/src/main/java/com/taobao/arthas/compiler/DynamicJavaFileManager.java b/memorycompiler/src/main/java/com/taobao/arthas/compiler/DynamicJavaFileManager.java
index 5ff09ff26..62d5d95e4 100644
--- a/memorycompiler/src/main/java/com/taobao/arthas/compiler/DynamicJavaFileManager.java
+++ b/memorycompiler/src/main/java/com/taobao/arthas/compiler/DynamicJavaFileManager.java
@@ -2,7 +2,7 @@ package com.taobao.arthas.compiler;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
 
@@ -13,6 +13,9 @@ import javax.tools.JavaFileObject;
 import javax.tools.StandardLocation;
 
 public class DynamicJavaFileManager extends ForwardingJavaFileManager<JavaFileManager> {
+    private static final String[] superLocationNames = { StandardLocation.PLATFORM_CLASS_PATH.name(),
+            /** JPMS StandardLocation.SYSTEM_MODULES **/
+            "SYSTEM_MODULES" };
     private final PackageInternalsFinder finder;
 
     private final DynamicClassLoader classLoader;
@@ -62,23 +65,62 @@ public class DynamicJavaFileManager extends ForwardingJavaFileManager<JavaFileMa
 
     @Override
     public Iterable<JavaFileObject> list(Location location, String packageName, Set<JavaFileObject.Kind> kinds,
-                    boolean recurse) throws IOException {
-        if (location == StandardLocation.PLATFORM_CLASS_PATH) { // let standard manager hanfle
-            return super.list(location, packageName, kinds, recurse);
-        } else if (location == StandardLocation.CLASS_PATH && kinds.contains(JavaFileObject.Kind.CLASS)) {
-            if (packageName.startsWith("java")) {
-                /**
-                 * a hack to let standard manager handle locations like "java.lang" or
-                 * "java.util". Prob would make sense to join results of standard manager with
-                 * those of my finder here
-                 */
-                return super.list(location, packageName, kinds, recurse);
-            } else { // app specific classes are here
-                return finder.find(packageName);
+                                         boolean recurse) throws IOException {
+        if (location instanceof StandardLocation) {
+            String locationName = ((StandardLocation) location).name();
+            for (String name : superLocationNames) {
+                if (name.equals(locationName)) {
+                    return super.list(location, packageName, kinds, recurse);
+                }
             }
         }
-        return Collections.emptyList();
 
+        // merge JavaFileObjects from specified ClassLoader
+        if (location == StandardLocation.CLASS_PATH && kinds.contains(JavaFileObject.Kind.CLASS)) {
+            return new IterableJoin<JavaFileObject>(super.list(location, packageName, kinds, recurse),
+                    finder.find(packageName));
+        }
+
+        return super.list(location, packageName, kinds, recurse);
     }
 
+    static class IterableJoin<T> implements Iterable<T> {
+        private final Iterable<T> first, next;
+
+        public IterableJoin(Iterable<T> first, Iterable<T> next) {
+            this.first = first;
+            this.next = next;
+        }
+
+        @Override
+        public Iterator<T> iterator() {
+            return new IteratorJoin<T>(first.iterator(), next.iterator());
+        }
+    }
+
+    static class IteratorJoin<T> implements Iterator<T> {
+        private final Iterator<T> first, next;
+
+        public IteratorJoin(Iterator<T> first, Iterator<T> next) {
+            this.first = first;
+            this.next = next;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return first.hasNext() || next.hasNext();
+        }
+
+        @Override
+        public T next() {
+            if (first.hasNext())
+                return first.next();
+            return next.next();
+        }
+
+        @Override
+        public void remove() {
+            throw new UnsupportedOperationException("remove");
+        }
+    }
 }