Add ability to filter out some particular classes from generator

compiler/backend/src/org/jetbrains/jet/codegen/PackageCodegen.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -174,8 +174,9 @@ public class PackageCodegen extends MemberCodegen {
                 generateSrcClass = true;
             }
             else if (declaration instanceof JetClassOrObject) {
-                if (state.isGenerateDeclaredClasses()) {
+                JetClassOrObject classOrObject = (JetClassOrObject) declaration;
-                    generateClassOrObject((JetClassOrObject) declaration);
+                if (state.getGenerateDeclaredClassFilter().shouldProcess(classOrObject)) {
+                    generateClassOrObject(classOrObject);
                 }
             }
             else if (declaration instanceof JetScript) {

compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenAnnotatingVisitor.java

@@ -24,6 +24,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.asm4.Type;
 import org.jetbrains.jet.codegen.SamCodegenUtil;
+import org.jetbrains.jet.codegen.state.GenerationState;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.impl.ClassDescriptorImpl;
 import org.jetbrains.jet.lang.psi.*;
@@ -53,7 +54,8 @@ import static org.jetbrains.jet.lexer.JetTokens.*;
 
 class CodegenAnnotatingVisitor extends JetVisitorVoid {
     private static final TokenSet BINARY_OPERATIONS = TokenSet.orSet(
-            AUGMENTED_ASSIGNMENTS, TokenSet.create(PLUS, MINUS, MUL, DIV, PERC, RANGE, LT, GT, LTEQ, GTEQ, IDENTIFIER));
+            AUGMENTED_ASSIGNMENTS,
+            TokenSet.create(PLUS, MINUS, MUL, DIV, PERC, RANGE, LT, GT, LTEQ, GTEQ, IDENTIFIER));
 
     private static class ClassDescriptorWithState {
 
@@ -82,11 +84,15 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
 
     private final Stack<ClassDescriptorWithState> classStack = new Stack<ClassDescriptorWithState>();
     private final Stack<String> nameStack = new Stack<String>();
+
     private final BindingTrace bindingTrace;
     private final BindingContext bindingContext;
 
-    public CodegenAnnotatingVisitor(BindingTrace bindingTrace) {
+    private final GenerationState.GenerateClassFilter filter;
+
+    public CodegenAnnotatingVisitor(BindingTrace bindingTrace, GenerationState.GenerateClassFilter filter) {
         this.bindingTrace = bindingTrace;
+        this.filter = filter;
         this.bindingContext = bindingTrace.getBindingContext();
     }
 
@@ -209,6 +215,8 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
 
     @Override
     public void visitClass(@NotNull JetClass klass) {
+        if (!filter.shouldProcess(klass)) return;
+
         ClassDescriptor classDescriptor = bindingContext.get(CLASS, klass);
         // working around a problem with shallow analysis
         if (classDescriptor == null) return;

compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenBinding.java

@@ -20,6 +20,7 @@ import com.intellij.openapi.vfs.VirtualFile;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.asm4.Type;
+import org.jetbrains.jet.codegen.state.GenerationState;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.impl.ClassDescriptorImpl;
 import org.jetbrains.jet.lang.psi.*;
@@ -60,7 +61,11 @@ public class CodegenBinding {
     }
 
     public static void initTrace(BindingTrace bindingTrace, Collection<JetFile> files) {
-        CodegenAnnotatingVisitor visitor = new CodegenAnnotatingVisitor(bindingTrace);
+        initTrace(bindingTrace, files, GenerationState.GenerateClassFilter.GENERATE_ALL);
+    }
+
+    public static void initTrace(BindingTrace bindingTrace, Collection<JetFile> files, GenerationState.GenerateClassFilter filter) {
+        CodegenAnnotatingVisitor visitor = new CodegenAnnotatingVisitor(bindingTrace, filter);
         for (JetFile file : allFilesInPackages(bindingTrace.getBindingContext(), files)) {
             file.accept(visitor);
         }
@@ -254,7 +259,7 @@ public class CodegenBinding {
             }
 
             @Override
-            public int compare(JetFile first, JetFile second) {
+            public int compare(@NotNull JetFile first, @NotNull JetFile second) {
                 return path(first).compareTo(path(second));
             }
         });

compiler/backend/src/org/jetbrains/jet/codegen/state/GenerationState.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@ import org.jetbrains.jet.codegen.SamWrapperClasses;
 import org.jetbrains.jet.codegen.binding.CodegenBinding;
 import org.jetbrains.jet.codegen.intrinsics.IntrinsicMethods;
 import org.jetbrains.jet.lang.descriptors.ScriptDescriptor;
+import org.jetbrains.jet.lang.psi.JetClassOrObject;
 import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.BindingTrace;
@@ -34,6 +35,24 @@ import org.jetbrains.jet.lang.resolve.DelegatingBindingTrace;
 import java.util.List;
 
 public class GenerationState {
+    public interface GenerateClassFilter {
+        boolean shouldProcess(JetClassOrObject classOrObject);
+
+        GenerateClassFilter ONLY_PACKAGE_CLASS = new GenerateClassFilter() {
+            @Override
+            public boolean shouldProcess(JetClassOrObject classOrObject) {
+                return false;
+            }
+        };
+
+        GenerateClassFilter GENERATE_ALL = new GenerateClassFilter() {
+            @Override
+            public boolean shouldProcess(JetClassOrObject classOrObject) {
+                return true;
+            }
+        };
+    }
+
     private boolean used = false;
 
     @NotNull
@@ -70,7 +89,7 @@ public class GenerationState {
 
     private final boolean generateNotNullParamAssertions;
 
-    private final boolean generateDeclaredClasses;
+    private final GenerateClassFilter generateClassFilter;
 
     private final boolean inlineEnabled;
 
@@ -84,7 +103,7 @@ public class GenerationState {
             @NotNull List<JetFile> files,
             boolean inlineEnabled
     ) {
-        this(project, builderFactory, Progress.DEAF, bindingContext, files, true, false, true, inlineEnabled);
+        this(project, builderFactory, Progress.DEAF, bindingContext, files, true, false, GenerateClassFilter.GENERATE_ALL, inlineEnabled);
     }
 
     public GenerationState(
@@ -95,7 +114,7 @@ public class GenerationState {
             @NotNull List<JetFile> files,
             boolean generateNotNullAssertions,
             boolean generateNotNullParamAssertions,
-            boolean generateDeclaredClasses,
+            GenerateClassFilter generateClassFilter,
             boolean inlineEnabled
     ) {
         this.project = project;
@@ -116,7 +135,7 @@ public class GenerationState {
 
         this.generateNotNullAssertions = generateNotNullAssertions;
         this.generateNotNullParamAssertions = generateNotNullParamAssertions;
-        this.generateDeclaredClasses = generateDeclaredClasses;
+        this.generateClassFilter = generateClassFilter;
     }
 
     @NotNull
@@ -177,8 +196,8 @@ public class GenerationState {
         return generateNotNullParamAssertions;
     }
 
-    public boolean isGenerateDeclaredClasses() {
+    public GenerateClassFilter getGenerateDeclaredClassFilter() {
-        return generateDeclaredClasses;
+        return generateClassFilter;
     }
 
     public boolean isInlineEnabled() {
@@ -189,7 +208,7 @@ public class GenerationState {
         markUsed();
 
         //noinspection unchecked
-        CodegenBinding.initTrace(getBindingTrace(), getFiles());
+        CodegenBinding.initTrace(getBindingTrace(), getFiles(), getGenerateDeclaredClassFilter());
     }
 
     private void markUsed() {

compiler/cli/src/org/jetbrains/jet/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -298,7 +298,7 @@ public class KotlinToJVMBytecodeCompiler {
                 project, ClassBuilderFactories.BINARIES, Progress.DEAF, exhaust.getBindingContext(), environment.getSourceFiles(),
                 configuration.get(JVMConfigurationKeys.GENERATE_NOT_NULL_ASSERTIONS, false),
                 configuration.get(JVMConfigurationKeys.GENERATE_NOT_NULL_PARAMETER_ASSERTIONS, false),
-                /*generateDeclaredClasses = */true,
+                GenerationState.GenerateClassFilter.GENERATE_ALL,
                 configuration.get(JVMConfigurationKeys.ENABLE_INLINE, InlineUtil.DEFAULT_INLINE_FLAG)
         );
         KotlinCodegenFacade.compileCorrectFiles(generationState, CompilationErrorHandler.THROW_EXCEPTION);

compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/KotlinJavaFileStubProvider.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,8 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.SystemInfo;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.ClassFileViewProvider;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiManager;
 import com.intellij.psi.impl.PsiManagerImpl;
 import com.intellij.psi.impl.compiled.ClsFileImpl;
@@ -33,6 +35,7 @@ import com.intellij.psi.stubs.PsiClassHolderFileStub;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.CachedValueProvider;
 import com.intellij.psi.util.PsiModificationTracker;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.testFramework.LightVirtualFile;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.Stack;
@@ -67,7 +70,7 @@ public class KotlinJavaFileStubProvider<T extends WithFileStub> implements Cache
         return new KotlinJavaFileStubProvider<KotlinPackageLightClassData>(
                 project,
                 false,
-                new StubGenerationStrategy.NoDeclaredClasses<KotlinPackageLightClassData>() {
+                new StubGenerationStrategy<KotlinPackageLightClassData>() {
                     @NotNull
                     @Override
                     public LightClassConstructionContext getContext(@NotNull Collection<JetFile> files) {
@@ -94,12 +97,22 @@ public class KotlinJavaFileStubProvider<T extends WithFileStub> implements Cache
                         return packageFqName;
                     }
 
+                    @Override
+                    public GenerationState.GenerateClassFilter getGenerateClassFilter() {
+                        return GenerationState.GenerateClassFilter.ONLY_PACKAGE_CLASS;
+                    }
+
                     @Override
                     public void generate(@NotNull GenerationState state, @NotNull Collection<JetFile> files) {
                         PackageCodegen codegen = state.getFactory().forPackage(packageFqName, files);
                         codegen.generate(CompilationErrorHandler.THROW_EXCEPTION);
                         state.getFactory().asList();
                     }
+
+                    @Override
+                    public String toString() {
+                        return StubGenerationStrategy.class.getName() + " for package class";
+                    }
                 }
         );
     }
@@ -109,7 +122,7 @@ public class KotlinJavaFileStubProvider<T extends WithFileStub> implements Cache
         return new KotlinJavaFileStubProvider<OutermostKotlinClassLightClassData>(
                 classOrObject.getProject(),
                 JetPsiUtil.isLocal(classOrObject),
-                new StubGenerationStrategy.WithDeclaredClasses<OutermostKotlinClassLightClassData>() {
+                new StubGenerationStrategy<OutermostKotlinClassLightClassData>() {
                     private JetFile getFile() {
                         return (JetFile) classOrObject.getContainingFile();
                     }
@@ -171,12 +184,55 @@ public class KotlinJavaFileStubProvider<T extends WithFileStub> implements Cache
                         return JetPsiUtil.getFQName(getFile());
                     }
 
+                    @Override
+                    public GenerationState.GenerateClassFilter getGenerateClassFilter() {
+                        return new GenerationState.GenerateClassFilter() {
+                            @Override
+                            public boolean shouldProcess(JetClassOrObject generatedClassOrObject) {
+                                // Trivial: generate and analyze class we are interested in.
+                                if (generatedClassOrObject == classOrObject) return true;
+
+                                // Process all parent classes as they are context for current class
+                                // Process child classes because they probably affect members (heuristic)
+                                if (PsiTreeUtil.isAncestor(generatedClassOrObject, classOrObject, true) ||
+                                    PsiTreeUtil.isAncestor(classOrObject, generatedClassOrObject, true)) {
+                                    return true;
+                                }
+
+                                if (JetPsiUtil.isLocal(generatedClassOrObject) && JetPsiUtil.isLocal(classOrObject)) {
+                                    // Local classes should be process by CodegenAnnotatingVisitor to
+                                    // decide what class they should be placed in.
+                                    //
+                                    // Example:
+                                    // class A
+                                    // fun foo() {
+                                    //     trait Z: A {}
+                                    //     fun bar() {
+                                    //         class <caret>O2: Z {}
+                                    //     }
+                                    // }
+
+                                    // TODO: current method will process local classes in irrelevant declarations, it should be fixed.
+                                    PsiElement commonParent = PsiTreeUtil.findCommonParent(generatedClassOrObject, classOrObject);
+                                    return commonParent != null && !(commonParent instanceof PsiFile);
+                                }
+
+                                return false;
+                            }
+                        };
+                    }
+
                     @Override
                     public void generate(@NotNull GenerationState state, @NotNull Collection<JetFile> files) {
                         PackageCodegen packageCodegen = state.getFactory().forPackage(getPackageFqName(), files);
                         packageCodegen.generateClassOrObject(classOrObject);
                         state.getFactory().asList();
                     }
+
+                    @Override
+                    public String toString() {
+                        return StubGenerationStrategy.class.getName() + " for explicit class " + classOrObject.getName();
+                    }
                 }
         );
     }
@@ -224,7 +280,7 @@ public class KotlinJavaFileStubProvider<T extends WithFileStub> implements Cache
                     context.getBindingContext(),
                     Lists.newArrayList(files),
                     /*not-null assertions*/false, false,
-                    /*generateDeclaredClasses=*/stubGenerationStrategy.generateDeclaredClasses(),
+                    /*generateClassFilter=*/stubGenerationStrategy.getGenerateClassFilter(),
                     InlineUtil.DEFAULT_INLINE_FLAG_FOR_STUB);
             state.beforeCompile();
 
@@ -309,36 +365,11 @@ public class KotlinJavaFileStubProvider<T extends WithFileStub> implements Cache
     private interface StubGenerationStrategy<T extends WithFileStub> {
         @NotNull Collection<JetFile> getFiles();
         @NotNull FqName getPackageFqName();
+
         @NotNull LightClassConstructionContext getContext(@NotNull Collection<JetFile> files);
         @NotNull T createLightClassData(PsiJavaFileStub javaFileStub, BindingContext bindingContext);
 
-        boolean generateDeclaredClasses();
+        GenerationState.GenerateClassFilter getGenerateClassFilter();
         void generate(@NotNull GenerationState state, @NotNull Collection<JetFile> files);
-
-        abstract class NoDeclaredClasses<U extends WithFileStub> implements StubGenerationStrategy<U> {
-            @Override
-            public boolean generateDeclaredClasses() {
-                return false;
-            }
-
-            @Override
-            public String toString() {
-                // For subclasses to be identifiable in the debugger
-                return NoDeclaredClasses.class.getSimpleName();
-            }
-        }
-
-        abstract class WithDeclaredClasses<U extends WithFileStub> implements StubGenerationStrategy<U> {
-            @Override
-            public boolean generateDeclaredClasses() {
-                return true;
-            }
-
-            @Override
-            public String toString() {
-                // For subclasses to be identifiable in the debugger
-                return WithDeclaredClasses.class.getSimpleName();
-            }
-        }
     }
 }

compiler/tests/org/jetbrains/jet/codegen/CodegenTestUtil.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -61,7 +61,7 @@ public class CodegenTestUtil {
                 environment.getProject(), ClassBuilderFactories.TEST, Progress.DEAF, analyzeExhaust.getBindingContext(), files.getPsiFiles(),
                 configuration.get(JVMConfigurationKeys.GENERATE_NOT_NULL_ASSERTIONS, true),
                 configuration.get(JVMConfigurationKeys.GENERATE_NOT_NULL_PARAMETER_ASSERTIONS, true),
-                /*generateDeclaredClasses = */true,
+                GenerationState.GenerateClassFilter.GENERATE_ALL,
                 configuration.get(JVMConfigurationKeys.ENABLE_INLINE, InlineUtil.DEFAULT_INLINE_FLAG_FOR_TEST)
         );
         KotlinCodegenFacade.compileCorrectFiles(state, CompilationErrorHandler.THROW_EXCEPTION);

idea/src/org/jetbrains/jet/plugin/internal/codewindow/BytecodeToolWindow.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -110,7 +110,8 @@ public class BytecodeToolWindow extends JPanel implements Disposable {
                     return printStackTraceToString(exhaust.getError());
                 }
                 state = new GenerationState(jetFile.getProject(), ClassBuilderFactories.TEXT, Progress.DEAF, exhaust.getBindingContext(),
-                                            Collections.singletonList(jetFile), true, true, true,
+                                            Collections.singletonList(jetFile), true, true,
+                                            GenerationState.GenerateClassFilter.GENERATE_ALL,
                                             InlineUtil.DEFAULT_INLINE_FLAG_FOR_TOOLWINDOW /*TODO add checkbox or extract it from option*/);
                 KotlinCodegenFacade.compileCorrectFiles(state, CompilationErrorHandler.THROW_EXCEPTION);
             }

idea/testData/javaFacade/localClassSubclass.kt

@@ -0,0 +1,7 @@
+class A
+fun foo() {
+    trait Z: A {}
+    fun bar() {
+        class <caret>O2: Z {}
+    }
+}

idea/tests/org/jetbrains/jet/plugin/javaFacade/JetJavaFacadeTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -59,6 +59,10 @@ public class JetJavaFacadeTest extends JetLightCodeInsightFixtureTestCase {
         doTestWrapMethod(true);
     }
 
+    public void testLocalClassSubclass() {
+        doTestWrapClass();
+    }
+
     public void testWrapFunInClassObject() {
         doTestWrapMethod(true);
     }