1package com.xtremelabs.robolectric.bytecode;
2
3import java.io.ByteArrayOutputStream;
4import java.io.File;
5import java.io.FileOutputStream;
6import java.io.IOException;
7import java.io.InputStream;
8import java.util.Enumeration;
9import java.util.HashMap;
10import java.util.Map;
11import java.util.jar.Attributes;
12import java.util.jar.JarEntry;
13import java.util.jar.JarFile;
14import java.util.jar.JarOutputStream;
15import java.util.jar.Manifest;
16
17public 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