From 7841961989165ee1983177a86eab72a42623057b Mon Sep 17 00:00:00 2001
From: Andrey Breslav <andrey.breslav@jetbrains.com>
Date: Tue, 23 Apr 2013 14:21:51 +0400
Subject: [PATCH] Callback added to preloader

---
 .../jet/preloading/ClassPreloadingUtils.java  | 27 ++++++++++++++++---
 .../jetbrains/jet/preloading/Preloader.java   | 27 ++++++++++++++-----
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/compiler/preloader/src/org/jetbrains/jet/preloading/ClassPreloadingUtils.java b/compiler/preloader/src/org/jetbrains/jet/preloading/ClassPreloadingUtils.java
index 7b9f02035e9..a7284f8bb45 100644
--- a/compiler/preloader/src/org/jetbrains/jet/preloading/ClassPreloadingUtils.java
+++ b/compiler/preloader/src/org/jetbrains/jet/preloading/ClassPreloadingUtils.java
@@ -29,6 +29,10 @@ import java.util.zip.ZipInputStream;
 
 public class ClassPreloadingUtils {
 
+    public interface ClassHandler {
+        void beforeDefineClass(String name, int sizeInBytes);
+    }
+
     /**
      * Creates a class loader that loads all classes from {@code jarFiles} into memory to make loading faster (avoid skipping through zip archives).
      *
@@ -37,18 +41,28 @@ public class ClassPreloadingUtils {
      * @param jarFiles jars to load all classes from
      * @param classCountEstimation an estimated number of classes in a the jars
      * @param parent (nullable) parent class loader
+     * @param handler handler to be notified on class definitions done by this class loader, or null
      * @return a class loader that reads classes from memory
      * @throws IOException on from reading the jar
      */
-    public static ClassLoader preloadClasses(Collection<File> jarFiles, int classCountEstimation, ClassLoader parent) throws IOException {
+    public static ClassLoader preloadClasses(
+            Collection<File> jarFiles, int classCountEstimation, ClassLoader parent, ClassHandler handler
+    ) throws IOException {
         Map<String, ResourceData> entries = loadAllClassesFromJars(jarFiles, classCountEstimation);
 
-        return createMemoryBasedClassLoader(parent, entries);
+        return createMemoryBasedClassLoader(parent, entries, handler);
+    }
+
+    public static ClassLoader preloadClasses(
+            Collection<File> jarFiles, int classCountEstimation, ClassLoader parent
+    ) throws IOException {
+        return preloadClasses(jarFiles, classCountEstimation, parent, null);
     }
 
     private static ClassLoader createMemoryBasedClassLoader(
             final ClassLoader parent,
-            final Map<String, ResourceData> preloadedResources
+            final Map<String, ResourceData> preloadedResources,
+            final ClassHandler handler
     ) {
         return new ClassLoader(null) {
             @Override
@@ -67,7 +81,12 @@ public class ClassPreloadingUtils {
                 ResourceData resourceData = preloadedResources.get(internalName);
                 if (resourceData == null) return null;
 
-                return defineClass(name, resourceData.bytes, 0, resourceData.bytes.length);
+                int sizeInBytes = resourceData.bytes.length;
+                if (handler != null) {
+                    handler.beforeDefineClass(name, sizeInBytes);
+                }
+
+                return defineClass(name, resourceData.bytes, 0, sizeInBytes);
             }
 
             @Override
diff --git a/compiler/preloader/src/org/jetbrains/jet/preloading/Preloader.java b/compiler/preloader/src/org/jetbrains/jet/preloading/Preloader.java
index 5be75667526..6c689834490 100644
--- a/compiler/preloader/src/org/jetbrains/jet/preloading/Preloader.java
+++ b/compiler/preloader/src/org/jetbrains/jet/preloading/Preloader.java
@@ -44,16 +44,31 @@ public class Preloader {
 
         ClassLoader parent = Preloader.class.getClassLoader();
 
-        ClassLoader preloaded = ClassPreloadingUtils.preloadClasses(files, classNumber, parent);
+        final int[] counter = new int[1];
+        final int[] size = new int[1];
+        ClassLoader preloaded = ClassPreloadingUtils.preloadClasses(files, classNumber, parent,
+                                                                    !printTime ? null :
+                                                                    new ClassPreloadingUtils.ClassHandler() {
+                                                                        @Override
+                                                                        public void beforeDefineClass(String name, int sizeInBytes) {
+                                                                            counter[0]++;
+                                                                            size[0] += sizeInBytes;
+                                                                        }
+                                                                    });
 
         Class<?> mainClass = preloaded.loadClass(mainClassCanonicalName);
         Method mainMethod = mainClass.getMethod("main", String[].class);
 
-        mainMethod.invoke(0, new Object[] {Arrays.copyOfRange(args, PRELOADER_ARG_COUNT, args.length)});
-
-        if (printTime) {
-            long dt = System.nanoTime() - startTime;
-            System.out.format("Total time: %.3fs\n", dt / 1e9);
+        try {
+            mainMethod.invoke(0, new Object[] {Arrays.copyOfRange(args, PRELOADER_ARG_COUNT, args.length)});
+        }
+        finally {
+            if (printTime) {
+                long dt = System.nanoTime() - startTime;
+                System.out.format("Total time: %.3fs\n", dt / 1e9);
+                System.out.println("Loaded classes: " + counter[0]);
+                System.out.println("Loaded classes size: " + size[0]);
+            }
         }
     }