querydsl/src/main/java/com/mysema/codegen/MemFileManager.java

/*
 * Copyright (c) 2010 Mysema Ltd.
 * 
 * base on code from https://hickory.dev.java.net/
 * 
 */

package com.mysema.codegen;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.StandardJavaFileManager;
import javax.tools.StandardLocation;
import javax.tools.JavaFileObject.Kind;



/**
 * MemFileManager is a memory based implementation of the JavaFileManager interface
 * 
 * @author tiwe
 *
 */
public class MemFileManager extends ForwardingJavaFileManager<JavaFileManager> {
    
    private final ClassLoader classLoader;
    
    private final Map<LocationAndKind, Map<String, JavaFileObject>> ramFileSystem;
    
    private final String urlPrefix;
    
    public MemFileManager(ClassLoader parent, StandardJavaFileManager sjfm) {
        super(sjfm);
        ramFileSystem = new HashMap<LocationAndKind,Map<String,JavaFileObject>>();
        Map<String,JavaFileObject> classLoaderContent = new HashMap<String,JavaFileObject>();
        ramFileSystem.put(new LocationAndKind(StandardLocation.CLASS_OUTPUT, Kind.CLASS),classLoaderContent);
        classLoader = new MemClassLoader(parent, ramFileSystem);
        urlPrefix = MemFileSystemRegistry.DEFAULT.getUrlPrefix(this);
    }
    
    @Override
    public ClassLoader getClassLoader(JavaFileManager.Location location) {
        return classLoader;
    }

    @Override
    public FileObject getFileForInput(JavaFileManager.Location location, String packageName, String relativeName) throws IOException {
        throw new UnsupportedOperationException();
    }
 
    @Override
    public FileObject getFileForOutput(JavaFileManager.Location location, String packageName, String relativeName, FileObject sibling) throws IOException {
        String name = null;
        if("".equals(packageName)) {
            name = relativeName;
        } else {
            name = packageName.replace('.','/') + "/" + relativeName;
        }
        LocationAndKind key = new LocationAndKind(location, Kind.OTHER);
        if(ramFileSystem.containsKey(key)) {
            JavaFileObject jfo = ramFileSystem.get(key).get(name);
            if(jfo != null) {
                return jfo;
            }
        }
        JavaFileObject jfo = new MemJavaFileObject(urlPrefix, name, Kind.OTHER);
        register(key,jfo);
        return jfo;
    }
    
    @Override
    public JavaFileObject getJavaFileForInput(JavaFileManager.Location location, String className, JavaFileObject.Kind kind) throws IOException {
        throw new UnsupportedOperationException();
    }

    @Override
    public JavaFileObject getJavaFileForOutput(Location location,
            String name, Kind kind, FileObject sibling)
            throws IOException {
        JavaFileObject javaFileObject = null;
        LocationAndKind key = new LocationAndKind(location,kind);

        if(ramFileSystem.containsKey(key)) {
            javaFileObject = ramFileSystem.get(key).get(name);
            if(javaFileObject != null) {
                return javaFileObject;
            }
        }
        if(kind == Kind.SOURCE) {
            javaFileObject = new MemSourceFileObject(name);
        } else {
            javaFileObject = new MemJavaFileObject(urlPrefix, name, kind);
        }  
        register(key,javaFileObject);
        return javaFileObject;
    }
    
    @Override
    public String inferBinaryName(Location loc, JavaFileObject javaFileObject) {
        String result;
        if (loc == StandardLocation.CLASS_PATH && javaFileObject instanceof MemJavaFileObject) {
            result = javaFileObject.getName();
        } else {
            result = super.inferBinaryName(loc, javaFileObject);
        }
        return result;
    }

    @Override
    public boolean isSameFile(FileObject a, FileObject b) {
        return a.equals(b);
    }
    
    @Override
    public Iterable<JavaFileObject> list(Location location, String pkg,
            Set<Kind> kinds, boolean recurse) throws IOException {
        
        List<JavaFileObject> result = new ArrayList<JavaFileObject>();
        for(JavaFileObject f : super.list(location, pkg, kinds, recurse)) {
            result.add(f);
        }
        if(location == StandardLocation.CLASS_PATH) {
            location = StandardLocation.CLASS_OUTPUT;
        }

        for(Kind kind : kinds) {
            LocationAndKind key = new LocationAndKind(location,kind);
            if(ramFileSystem.containsKey(key)) {
                Map<String, JavaFileObject> locatedFiles = ramFileSystem.get(key);
                for (Map.Entry<String,JavaFileObject> entry : locatedFiles.entrySet()){                    
                    String name = entry.getKey();
                    String packageName = "";
                    if (name.indexOf('.') > -1){
                        packageName = name.substring(0,name.lastIndexOf('.'));
                    }
                    if(recurse ? packageName.startsWith(pkg) : packageName.equals(pkg)) {
                        JavaFileObject candidate = entry.getValue();
                        if(kinds.contains(candidate.getKind())) {
                            result.add(candidate);
                        }
                    }
                }
            }
        }
        return result;
    }

    private void register(LocationAndKind key, JavaFileObject javaFileObject) {
        if(!ramFileSystem.containsKey(key)) {
            ramFileSystem.put(key, new HashMap<String, JavaFileObject>());
        }
        ramFileSystem.get(key).put(javaFileObject.getName(), javaFileObject);
    }

}