Home | History | Annotate | Download | only in bytecode 
      1 package com.xtremelabs.robolectric.bytecode;
      2 
      3 import java.io.ByteArrayOutputStream;
      4 import java.io.File;
      5 import java.io.FileOutputStream;
      6 import java.io.IOException;
      7 import java.io.InputStream;
      8 import java.util.Enumeration;
      9 import java.util.HashMap;
     10 import java.util.Map;
     11 import java.util.jar.Attributes;
     12 import java.util.jar.JarEntry;
     13 import java.util.jar.JarFile;
     14 import java.util.jar.JarOutputStream;
     15 import java.util.jar.Manifest;
     16 
     17 public class ClassCache {
     18     private static final Attributes.Name VERSION_ATTRIBUTE = new Attributes.Name("version");
     19 
     20     private Map<String, byte[]> cachedClasses = new HashMap<String, byte[]>();
     21     private boolean startedWriting = false;
     22 
     23     public ClassCache(String classCachePath, final int expectedCacheVersion) {
     24         final File cacheJarFile = new File(classCachePath);
     25         JarFile cacheFile;
     26         try {
     27             cacheFile = new JarFile(cacheJarFile);
     28             int cacheVersion = 0;
     29             Manifest manifest = cacheFile.getManifest();
     30             if (manifest != null) {
     31                 Attributes attributes = manifest.getEntries().get("robolectric");
     32                 if (attributes != null) {
     33                     String cacheVersionStr = (String) attributes.get(VERSION_ATTRIBUTE);
     34                     if (cacheVersionStr != null) {
     35                         cacheVersion = Integer.parseInt(cacheVersionStr);
     36                     }
     37                 }
     38             }
     39             if (cacheVersion != expectedCacheVersion) {
     40                 cacheJarFile.delete();
     41             } else {
     42                 readEntries(cacheFile);
     43             }
     44         } catch (IOException e) {
     45             // no problem
     46         }
     47 
     48         Runtime.getRuntime().addShutdownHook(new Thread() {
     49             @Override public void run() {
     50                 Manifest manifest = new Manifest();
     51                 Attributes attributes = new Attributes();
     52                 attributes.put(VERSION_ATTRIBUTE, String.valueOf(expectedCacheVersion));
     53                 manifest.getEntries().put("robolectric", attributes);
     54 
     55                 saveAllClassesToCache(cacheJarFile, manifest);
     56             }
     57         });
     58     }
     59 
     60     public byte[] getClassBytesFor(String name) {
     61         return cachedClasses.get(name);
     62     }
     63 
     64     public boolean isWriting() {
     65         synchronized (this) {
     66             return startedWriting;
     67         }
     68     }
     69 
     70     public void addClass(String className, byte[] classBytes) {
     71         cachedClasses.put(className, classBytes);
     72     }
     73 
     74     private void readEntries(JarFile cacheFile) {
     75         Enumeration<JarEntry> entries = cacheFile.entries();
     76         try {
     77             while (entries.hasMoreElements()) {
     78                 JarEntry entry = entries.nextElement();
     79                 String className = entry.getName();
     80                 if (className.endsWith(".class")) {
     81                     int classSize = (int) entry.getSize();
     82                     InputStream inputStream = cacheFile.getInputStream(entry);
     83                     ByteArrayOutputStream baos = new ByteArrayOutputStream(classSize);
     84                     int c;
     85                     while ((c = inputStream.read()) != -1) {
     86                         baos.write(c);
     87                     }
     88                     className = className.substring(0, className.indexOf(".class")).replace('/', '.');
     89                     addClass(className, baos.toByteArray());
     90                 }
     91 
     92             }
     93         } catch (IOException e) {
     94             // no problem, we didn't want those bytes that much anyway
     95         }
     96     }
     97 
     98     protected void saveAllClassesToCache(File file, Manifest manifest) {
     99         synchronized (this) {
    100             startedWriting = true;
    101 
    102             if (cachedClasses.size() > 0) {
    103                 JarOutputStream jarOutputStream = null;
    104                 try {
    105                     File cacheJarDir = file.getParentFile();
    106                     if (!cacheJarDir.exists()) {
    107                         cacheJarDir.mkdirs();
    108                     }
    109 
    110                     jarOutputStream = new JarOutputStream(new FileOutputStream(file), manifest);
    111                     for (Map.Entry<String, byte[]> entry : cachedClasses.entrySet()) {
    112                         String key = entry.getKey();
    113                         jarOutputStream.putNextEntry(new JarEntry(key.replace('.', '/') + ".class"));
    114                         jarOutputStream.write(entry.getValue());
    115                         jarOutputStream.closeEntry();
    116                     }
    117                 } catch (IOException e) {
    118                     throw new RuntimeException(e);
    119                 } finally {
    120                     if (jarOutputStream != null) {
    121                         try {
    122                             jarOutputStream.close();
    123                         } catch (IOException ignore) {
    124                         }
    125                     }
    126                 }
    127             }
    128             startedWriting = false;
    129         }
    130     }
    131 }
    132