001    /*
002            Copyright (c) 2009 Olivier Chafik, All Rights Reserved
003            
004            This file is part of JNAerator (http://jnaerator.googlecode.com/).
005            
006            JNAerator is free software: you can redistribute it and/or modify
007            it under the terms of the GNU General Public License as published by
008            the Free Software Foundation, either version 3 of the License, or
009            (at your option) any later version.
010            
011            JNAerator is distributed in the hope that it will be useful,
012            but WITHOUT ANY WARRANTY; without even the implied warranty of
013            MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
014            GNU General Public License for more details.
015            
016            You should have received a copy of the GNU General Public License
017            along with JNAerator.  If not, see <http://www.gnu.org/licenses/>.
018    */
019    package com.ochafik.lang.compiler;
020    
021    import java.io.File;
022    import java.io.FileInputStream;
023    import java.io.FileNotFoundException;
024    import java.io.FileOutputStream;
025    import java.io.IOException;
026    import java.io.InputStream;
027    import java.io.OutputStream;
028    import java.util.ArrayList;
029    import java.util.HashMap;
030    import java.util.Map;
031    import java.util.jar.JarEntry;
032    import java.util.jar.JarOutputStream;
033    
034    import javax.tools.FileObject;
035    import javax.tools.ForwardingJavaFileManager;
036    import javax.tools.JavaFileManager;
037    import javax.tools.JavaFileObject;
038    
039    import com.ochafik.io.IOUtils;
040    
041    
042    public class MemoryFileManager extends ForwardingJavaFileManager<JavaFileManager> {
043            public final Map<String, MemoryJavaFile> inputs = new HashMap<String, MemoryJavaFile>();
044            public final Map<String, FileObject> outputs = new HashMap<String, FileObject>();
045    
046            public void writeJar(File outputJar, boolean outputSources, Map<String, File> additionalFiles) throws IOException {
047    
048                    File jarDir = outputJar.getAbsoluteFile().getParentFile();
049                    if (!jarDir.isDirectory())
050                            jarDir.mkdirs();
051                    writeJar(new FileOutputStream(outputJar), outputSources, additionalFiles);
052            }
053            public void writeJar(OutputStream out, boolean outputSources, Map<String, File> additionalFiles) throws IOException {
054                    JarOutputStream jout = new JarOutputStream(out);
055                    if (outputSources)
056                            for (Map.Entry<String, MemoryJavaFile> e : inputs.entrySet()) {
057    //                              if (e.getKey().contains(".class"))
058    //                                      continue;
059                                    writeEntry(e.getKey(), e.getValue(), jout);
060                            }
061                    for (Map.Entry<String, FileObject> e : outputs.entrySet())
062                            writeEntry(e.getKey(), e.getValue(), jout);
063                    
064                    if (additionalFiles != null)
065                            for (Map.Entry<String, File> additionalFile : additionalFiles.entrySet()) {
066                                    String path = additionalFile.getKey();
067                                    if (path.startsWith("file:///"))
068                                            path = path.substring("file:///".length());
069    
070                                    FileInputStream in = new FileInputStream(additionalFile.getValue());
071                                    JarEntry e = new JarEntry(path);
072                                    jout.putNextEntry(e);
073                                    IOUtils.readWrite(in, jout);
074                                    in.close();
075                                    jout.closeEntry();
076                            }
077                    jout.close();
078            }
079            protected void writeEntry(String path, FileObject o, JarOutputStream jout) throws IOException {
080                    if (path.startsWith("/"))
081                            path = path.substring(1);
082                    if (path.startsWith("file:///"))
083                            path = path.substring("file:///".length());
084    
085                    if (o instanceof MemoryFileObject) {
086                            MemoryFileObject mo = (MemoryFileObject)o;
087                            byte[] c = mo.getContent();
088                            if (c == null)
089                                    return;
090            
091    //                      String path = mo.getPath();
092    //                      if (path.startsWith("file:///"))
093    //                              path = path.substring("file:///".length());
094                            JarEntry e = new JarEntry(path);
095                            jout.putNextEntry(e);
096                            jout.write(c);
097                            jout.closeEntry();
098                    } else if (o instanceof URLFileObject) {
099                            URLFileObject uo = (URLFileObject)o;
100    //                      String path = uo.url.getFile();
101    //                      if (path.startsWith("/"))
102    //                              path = path.substring(1);
103                            
104                            JarEntry e = new JarEntry(path);
105                            jout.putNextEntry(e);
106                            InputStream in = uo.url.openStream();
107                            if (in == null)
108                                    throw new FileNotFoundException(path);
109                            IOUtils.readWrite(in, jout);
110                            in.close();
111                            jout.closeEntry();
112                    } else
113                            throw new UnsupportedOperationException("Dunno how to deal with " + o);
114            }
115            public MemoryFileManager(JavaFileManager fm) {
116                    super(fm);
117            }
118    
119            public MemoryJavaFile addSourceInput(String path, String content) {
120                    if (!path.startsWith("file:///"))
121                            path = "file:///" + path;
122                    
123                    MemoryJavaFile mjf = new MemoryJavaFile(path, content, JavaFileObject.Kind.SOURCE);
124                    inputs.put(path, mjf);
125                    return mjf;
126            }
127            @Override
128            public boolean isSameFile(FileObject a, FileObject b) {
129                    return a.toString().equals(b.toString());
130            }
131            public Iterable<? extends JavaFileObject> getJavaFileObjects() {
132                    return new ArrayList<JavaFileObject>(inputs.values());
133            }
134    
135            @Override
136            public JavaFileObject getJavaFileForInput(Location location, String className, JavaFileObject.Kind kind) throws IOException {
137    //              System.out.println("getJavaFileForInput(className = " + className + ", location = " + location + ", kind = " + kind + ")");
138                    if (kind == JavaFileObject.Kind.SOURCE) {
139                            return inputs.get(className);
140                    }
141                    return super.getJavaFileForInput(location, className, kind);
142            }
143    
144            public static String getFullPathForClass(String className, String extension) {
145                    return "file:///" + getSimplePathForClass(className, extension);
146            }
147            public static String getSimplePathForClass(String className, String extension) {
148                    return className.replace('.', '/') + "." + extension;
149            }
150            @Override
151            public JavaFileObject getJavaFileForOutput(Location location, String className, JavaFileObject.Kind kind, FileObject sibling) throws IOException {
152    //              System.out.println("getJavaFileForOutput(className = " + className + ", location = " + location + ", kind = " + kind + ")");
153                    MemoryJavaFile jo = null;
154                    if (kind == JavaFileObject.Kind.CLASS) {
155                            outputs.put(getFullPathForClass(className, "class"), jo = new MemoryJavaFile(getFullPathForClass(className, "class"), null, kind));
156                    } else if (kind == JavaFileObject.Kind.SOURCE) {
157                            inputs.put(getFullPathForClass(className, "java"), jo = new MemoryJavaFile(getFullPathForClass(className, "java"), null, kind));
158                    }
159    
160                    return jo == null ? super.getJavaFileForInput(location, className, kind) : jo;
161            }
162            @Override
163            public FileObject getFileForOutput(Location location, String packageName, String relativeName, FileObject sibling) throws IOException {
164    //              System.out.println("getFileForOutput(relativeName = " + relativeName + ")");
165                    if (relativeName.startsWith("file:///"))
166                            relativeName = relativeName.substring("file:///".length());
167                    
168                    FileObject out = outputs.get(relativeName);
169                    if (out == null) {
170                            out = new MemoryFileObject(relativeName, (String)null);
171                            outputs.put(relativeName, out);
172                    }
173                    return out;
174            }
175    }