PsiJavaStubModelRoot.java

package jetbrains.mps.idea.java.psiStubs;

/*Generated by MPS */

import jetbrains.mps.extapi.persistence.ModelRootBase;
import jetbrains.mps.idea.java.psi.JavaPsiListener;
import org.apache.log4j.Logger;
import org.jetbrains.annotations.NotNull;
import com.intellij.openapi.module.Module;
import org.jetbrains.mps.openapi.model.SModel;
import org.jetbrains.mps.openapi.model.SModelId;
import jetbrains.mps.idea.java.psi.PsiChangesWatcher;
import java.util.Map;
import org.jetbrains.mps.openapi.model.SModelReference;
import java.util.List;
import com.intellij.psi.PsiDirectory;
import jetbrains.mps.internal.collections.runtime.MapSequence;
import java.util.HashMap;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.openapi.roots.ModuleRootManager;
import com.intellij.psi.PsiManager;
import jetbrains.mps.internal.collections.runtime.ListSequence;
import java.util.ArrayList;
import jetbrains.mps.internal.collections.runtime.IMapping;
import jetbrains.mps.internal.collections.runtime.Sequence;
import com.intellij.psi.PsiJavaFile;
import jetbrains.mps.java.stub.JavaPackageNameStub;
import org.jetbrains.mps.openapi.persistence.Memento;
import com.intellij.psi.PsiFileSystemItem;

public class PsiJavaStubModelRoot extends ModelRootBase implements JavaPsiListener {
  private static Logger LOG = Logger.getLogger(PsiJavaStubModelRoot.class);
  private static final String TYPE = "JavaPsiStubs";

  @NotNull
  private Module myIdeaModule;

  public PsiJavaStubModelRoot(Module module) {
    myIdeaModule = module;
  }

  /**
   * Equals is defined only by our ideaModule, all the state is not taken into acount
   * We should be careful when working with ModelRoots in collections (see AbstractModule.doUpdateModelSet)
   */
  @Override
  public boolean equals(Object root) {
    if (!(root instanceof PsiJavaStubModelRoot)) {
      return false;
    }
    return myIdeaModule.equals(((PsiJavaStubModelRoot) root).myIdeaModule);
  }
  @Override
  public int hashCode() {
    return myIdeaModule.hashCode();
  }
  @Override
  public String getType() {
    return TYPE;
  }
  @Override
  public String getPresentation() {
    return "Java PSI stubs";
  }
  @Override
  public SModel getModel(SModelId id) {
    // TODO 
    return null;
  }
  @Override
  public void attach() {
    super.attach();

    //  start to listen 
    PsiChangesWatcher w = myIdeaModule.getProject().getComponent(PsiChangesWatcher.class);
    w.addListener(this);
  }
  @Override
  public void dispose() {
    super.dispose();
    PsiChangesWatcher w = myIdeaModule.getProject().getComponent(PsiChangesWatcher.class);
    w.removeListener(this);
  }
  @Override
  public Iterable<SModel> loadModels() {
    Map<SModelReference, List<PsiDirectory>> packageToDirs = MapSequence.fromMap(new HashMap<SModelReference, List<PsiDirectory>>());

    final VirtualFile[] sourceRoots = ModuleRootManager.getInstance(myIdeaModule).getSourceRoots(false);
    PsiManager psiMgr = PsiManager.getInstance(myIdeaModule.getProject());

    for (VirtualFile root : sourceRoots) {
      PsiDirectory dir = psiMgr.findDirectory(root);
      // judging by RootModelBase.getSourceRoots() only valid source roots will be returned, but we'll be paranoid 
      if (dir == null) {
        continue;
      }
      collectPackagesInDir(dir, dir, packageToDirs);
    }

    List<SModel> models = ListSequence.fromList(new ArrayList<SModel>());

    for (IMapping<SModelReference, List<PsiDirectory>> pair : MapSequence.fromMap(packageToDirs)) {
      SModelReference modelRef = pair.key();
      List<PsiDirectory> dirs = pair.value();

      ListSequence.fromList(models).addElement(makeModelDescriptor(modelRef, dirs));
    }

    return models;
  }
  private void collectPackagesInDir(PsiDirectory sourceRoot, PsiDirectory dir, Map<SModelReference, List<PsiDirectory>> result) {
    if (Sequence.fromIterable(Sequence.fromArray(dir.getFiles())).ofType(PsiJavaFile.class).isNotEmpty()) {

      SModelReference modelRef = makeModelReference(sourceRoot, dir);
      if (modelRef != null) {
        List<PsiDirectory> dirs = MapSequence.fromMap(result).get(modelRef);
        if (dirs == null) {
          dirs = ListSequence.fromList(new ArrayList<PsiDirectory>());
          MapSequence.fromMap(result).put(modelRef, dirs);
        }
        ListSequence.fromList(dirs).addElement(dir);
      }
    }

    for (PsiDirectory subDir : dir.getSubdirectories()) {
      collectPackagesInDir(sourceRoot, subDir, result);
    }
  }
  private PsiJavaStubModelDescriptor makeModelDescriptor(SModelReference modelRef, Iterable<PsiDirectory> dirs) {
    MultiplePsiJavaStubDataSource ds = new MultiplePsiJavaStubDataSource(myIdeaModule, dirs);
    PsiJavaStubModelDescriptor md = new PsiJavaStubModelDescriptor(modelRef, ds);
    md.setModelRoot(this);
    return md;
  }
  private SModelReference makeModelReference(PsiDirectory sourceRoot, PsiDirectory dir) {
    int skipPrefix = sourceRoot.toString().length();
    String relativeDirName = dir.toString().substring(skipPrefix);
    String packageName = relativeDirName.replace('/', '.').replace('\\', '.');

    if ((packageName == null || packageName.length() == 0)) {
      return null;
    }

    if (packageName.length() > 0 && packageName.charAt(0) == '.') {
      packageName = packageName.substring(1);
    }

    return new JavaPackageNameStub(packageName).asModelReference(getModule().getModuleReference());
  }
  public boolean isReadOnly() {
    return true;
  }
  @Override
  public boolean canCreateModel(String modelName) {
    return false;
  }
  @Override
  public SModel createModel(String modelName) {
    return null;
  }
  @Override
  public void save(Memento memento) {
    throw new UnsupportedOperationException("JavaPsiStubs: unsupported for now");
  }
  @Override
  public void load(Memento memento) {
    throw new UnsupportedOperationException("JavaPsiStubs: unsupported for now");
  }

  @Override
  public void psiChanged(JavaPsiListener.PsiEvent event) {

    // TODO re-write to sequences and any 

    // here we simply decide if we have to update 
    for (PsiFileSystemItem fsItem : event.getCreated()) {
      if (importantFsItem(fsItem)) {
        update();
        return;
      }
    }

    for (PsiFileSystemItem fsItem : event.getRemoved()) {
      if (importantFsItem(fsItem)) {
        update();
        return;
      }
    }

    for (JavaPsiListener.FSRename rename : event.getRenamed()) {
      if (importantFsItem(rename.item)) {
        update();
        return;
      }
    }

    for (JavaPsiListener.FSMove move : event.getMoved()) {
      if (importantFsItem(move.item)) {
        update();
        return;
      }
    }
  }


  private boolean importantFsItem(PsiFileSystemItem fsItem) {
    return (fsItem instanceof PsiDirectory || fsItem instanceof PsiJavaFile) && findOurSourceRoot(fsItem) != null;
  }

  private PsiDirectory findOurSourceRoot(PsiFileSystemItem item) {
    for (VirtualFile sourceRoot : ModuleRootManager.getInstance(myIdeaModule).getSourceRoots()) {
      String rootPath = sourceRoot.toString();
      String itemPath = item.getVirtualFile().toString();
      if (itemPath.startsWith(rootPath)) {
        return PsiManager.getInstance(myIdeaModule.getProject()).findDirectory(sourceRoot);
      }
    }
    return null;
  }
}