[groovy] refactor support of generics in compiled traits
[idea/community.git] / plugins / groovy / groovy-psi / src / org / jetbrains / plugins / groovy / lang / psi / util / GrTraitUtil.java
1 /*
2  * Copyright 2000-2015 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package org.jetbrains.plugins.groovy.lang.psi.util;
17
18 import com.intellij.codeInsight.AnnotationUtil;
19 import com.intellij.openapi.diagnostic.Logger;
20 import com.intellij.openapi.util.Ref;
21 import com.intellij.psi.*;
22 import com.intellij.psi.impl.compiled.ClsClassImpl;
23 import com.intellij.util.Function;
24 import com.intellij.util.containers.ContainerUtil;
25 import org.jetbrains.annotations.Contract;
26 import org.jetbrains.annotations.NotNull;
27 import org.jetbrains.annotations.Nullable;
28 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.GrTypeDefinition;
29 import org.jetbrains.plugins.groovy.lang.psi.impl.synthetic.GrLightMethodBuilder;
30
31 import java.util.Map;
32
33 import static com.intellij.psi.PsiModifier.ABSTRACT;
34
35 /**
36  * Created by Max Medvedev on 16/05/14
37  */
38 public class GrTraitUtil {
39   private static final Logger LOG = Logger.getInstance(GrTraitUtil.class);
40   private static final PsiTypeMapper ID_MAPPER = new PsiTypeMapper() {
41     @Override
42     public PsiType visitClassType(PsiClassType classType) {
43       return classType;
44     }
45   };
46
47   @Contract("null -> false")
48   public static boolean isInterface(@Nullable PsiClass aClass) {
49     return aClass != null && aClass.isInterface() && !isTrait(aClass);
50   }
51
52   public static boolean isMethodAbstract(PsiMethod method) {
53     if (method.getModifierList().hasExplicitModifier(ABSTRACT)) return true;
54
55     PsiClass aClass = method.getContainingClass();
56     return isInterface(aClass);
57   }
58
59   @NotNull
60   public static String getTraitFieldPrefix(@NotNull PsiClass aClass) {
61     String qname = aClass.getQualifiedName();
62     LOG.assertTrue(qname != null, aClass.getClass());
63
64     String[] idents = qname.split("\\.");
65
66     StringBuilder buffer = new StringBuilder();
67     for (String ident : idents) {
68       buffer.append(ident).append("_");
69     }
70
71     buffer.append("_");
72     return buffer.toString();
73   }
74
75   @Contract("null -> false")
76   public static boolean isTrait(@Nullable PsiClass containingClass) {
77     return containingClass instanceof GrTypeDefinition && ((GrTypeDefinition)containingClass).isTrait()
78            || containingClass instanceof ClsClassImpl
79               && containingClass.isInterface()
80               && AnnotationUtil.isAnnotated(containingClass, "groovy.transform.Trait", false);
81   }
82
83   public static PsiMethod createTraitMethodFromCompiledHelperMethod(final PsiMethod compiledMethod, final PsiClass trait) {
84     assert compiledMethod.getParameterList().getParametersCount() > 0;
85
86     final GrLightMethodBuilder result = new GrLightMethodBuilder(compiledMethod.getManager(), compiledMethod.getName());
87     result.setNavigationElement(compiledMethod);
88     result.setOriginInfo("via @Trait");
89     result.addModifier(PsiModifier.STATIC);
90     for (PsiTypeParameter parameter : compiledMethod.getTypeParameters()) {
91       result.getTypeParameterList().addParameter(parameter);
92     }
93
94     final PsiTypeVisitor<PsiType> corrector = createCorrector(compiledMethod, trait);
95
96     final PsiParameter[] methodParameters = compiledMethod.getParameterList().getParameters();
97     for (int i = 1; i < methodParameters.length; i++) {
98       final PsiParameter originalParameter = methodParameters[i];
99       final PsiType correctedType = originalParameter.getType().accept(corrector);
100       result.addParameter(originalParameter.getName(), correctedType, false);
101     }
102
103     for (PsiClassType type : compiledMethod.getThrowsList().getReferencedTypes()) {
104       final PsiType correctedType = type.accept(corrector);
105       result.getThrowsList().addReference(correctedType instanceof PsiClassType ? (PsiClassType)correctedType : type);
106     }
107
108     {
109       final PsiType originalType = compiledMethod.getReturnType();
110       result.setReturnType(originalType == null ? null : originalType.accept(corrector));
111     }
112
113     return result;
114   }
115
116   @NotNull
117   private static PsiTypeMapper createCorrector(final PsiMethod compiledMethod, final PsiClass trait) {
118     final PsiTypeParameter[] traitTypeParameters = trait.getTypeParameters();
119     if (traitTypeParameters.length == 0) return ID_MAPPER;
120
121     final Map<String, PsiTypeParameter> substitutionMap = ContainerUtil.newTroveMap();
122     for (PsiTypeParameter parameter : traitTypeParameters) {
123       substitutionMap.put(parameter.getName(), parameter);
124     }
125
126     final PsiElementFactory elementFactory = JavaPsiFacade.getInstance(trait.getProject()).getElementFactory();
127     return new PsiTypeMapper() {
128
129       @Nullable
130       @Override
131       public PsiType visitClassType(PsiClassType originalType) {
132         final PsiClass resolved = originalType.resolve();
133         // if resolved to method parameter -> return as is
134         if (resolved instanceof PsiTypeParameter && compiledMethod.equals(((PsiTypeParameter)resolved).getOwner())) return originalType;
135         final PsiType[] typeParameters = originalType.getParameters();
136         final PsiTypeParameter byName = substitutionMap.get(originalType.getCanonicalText());
137         if (byName != null) {
138           assert typeParameters.length == 0;
139           return elementFactory.createType(byName);
140         }
141         if (resolved == null) return originalType;
142         if (typeParameters.length == 0) return originalType;    // do not go deeper
143
144         final Ref<Boolean> hasChanges = Ref.create(false);
145         final PsiTypeVisitor<PsiType> $this = this;
146         final PsiType[] substitutes = ContainerUtil.map2Array(typeParameters, PsiType.class, new Function<PsiType, PsiType>() {
147           @Override
148           public PsiType fun(PsiType type) {
149             final PsiType mapped = type.accept($this);
150             hasChanges.set(mapped != type);
151             return mapped;
152           }
153         });
154         return hasChanges.get() ? elementFactory.createType(resolved, substitutes) : originalType;
155       }
156     };
157   }
158 }