Cache results in AST instead of a map

javaparser-core/src/main/java/com/github/javaparser/ast/Node.java

@@ -39,6 +39,7 @@
 import com.github.javaparser.printer.PrettyPrinter;
 import com.github.javaparser.printer.PrettyPrinterConfiguration;
 import com.github.javaparser.resolution.SymbolResolver;
+import com.github.javaparser.resolution.types.ResolvedType;
 
 import java.util.*;
 import java.util.function.Consumer;
@@ -53,10 +54,6 @@
 import static java.util.Spliterator.DISTINCT;
 import static java.util.Spliterator.NONNULL;
 
-import com.github.javaparser.metamodel.NodeMetaModel;
-import com.github.javaparser.metamodel.JavaParserMetaModel;
-import com.github.javaparser.ast.Node;
-
 /**
  * Base class for all nodes of the abstract syntax tree.
  * <h2>Construction</h2>
@@ -477,12 +474,23 @@ public <M> void setData(DataKey<M> key, M object) {
 
     /**
      * @return does this node have data for this key?
+     * @see DataKey
      */
     public boolean containsData(DataKey<?> key) {
         if (data == null) {
             return false;
         }
-        return data.get(key) != null;
+        return data.containsKey(key);
+    }
+
+    /**
+     * Remove data by key.
+     * @see DataKey
+     */
+    public void removeData(DataKey<ResolvedType> key){
+        if(data!=null){
+            data.remove(key);
+        }
     }
 
     /**

javaparser-symbol-solver-core/src/main/java/com/github/javaparser/symbolsolver/javaparsermodel/JavaParserFacade.java

@@ -17,10 +17,13 @@
 package com.github.javaparser.symbolsolver.javaparsermodel;
 
 import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.DataKey;
 import com.github.javaparser.ast.Node;
 import com.github.javaparser.ast.NodeList;
-import com.github.javaparser.ast.body.*;
+import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
 import com.github.javaparser.ast.body.EnumDeclaration;
+import com.github.javaparser.ast.body.TypeDeclaration;
+import com.github.javaparser.ast.body.VariableDeclarator;
 import com.github.javaparser.ast.expr.*;
 import com.github.javaparser.ast.stmt.ExplicitConstructorInvocationStmt;
 import com.github.javaparser.ast.type.*;
@@ -30,10 +33,12 @@
 import com.github.javaparser.resolution.types.*;
 import com.github.javaparser.symbolsolver.core.resolution.Context;
 import com.github.javaparser.symbolsolver.javaparsermodel.contexts.FieldAccessContext;
-import com.github.javaparser.symbolsolver.javaparsermodel.declarations.*;
+import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserAnonymousClassDeclaration;
+import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserEnumDeclaration;
+import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserTypeVariableDeclaration;
 import com.github.javaparser.symbolsolver.model.resolution.SymbolReference;
 import com.github.javaparser.symbolsolver.model.resolution.TypeSolver;
-import com.github.javaparser.symbolsolver.model.typesystem.*;
+import com.github.javaparser.symbolsolver.model.typesystem.ReferenceTypeImpl;
 import com.github.javaparser.symbolsolver.reflectionmodel.ReflectionClassDeclaration;
 import com.github.javaparser.symbolsolver.resolution.ConstructorResolutionLogic;
 import com.github.javaparser.symbolsolver.resolution.SymbolSolver;
@@ -51,12 +56,15 @@
  */
 public class JavaParserFacade {
 
-    private static Map<TypeSolver, JavaParserFacade> instances = new WeakHashMap<>();
-    private TypeSolver typeSolver;
-    private SymbolSolver symbolSolver;
-    private Map<Node, ResolvedType> cacheWithLambdasSolved = new IdentityHashMap<>();
-    private Map<Node, ResolvedType> cacheWithoutLambdasSolved = new IdentityHashMap<>();
-    private TypeExtractor typeExtractor;
+    private static final DataKey<ResolvedType> TYPE_WITH_LAMBDAS_RESOLVED = new DataKey<ResolvedType>() {
+    };
+    private static final DataKey<ResolvedType> TYPE_WITHOUT_LAMBDAS_RESOLVED = new DataKey<ResolvedType>() {
+    };
+
+    private static final Map<TypeSolver, JavaParserFacade> instances = new WeakHashMap<>();
+    private final TypeSolver typeSolver;
+    private final TypeExtractor typeExtractor;
+    private final SymbolSolver symbolSolver;
 
     private JavaParserFacade(TypeSolver typeSolver) {
         this.typeSolver = typeSolver.getRoot();
@@ -255,27 +263,27 @@ public SymbolReference<ResolvedValueDeclaration> solve(FieldAccessExpr fieldAcce
 
     /**
      * Get the type associated with the node.
-     *
+     * <p>
      * This method was originally intended to get the type of a value: any value has a type.
-     *
+     * <p>
      * For example:
      * <code>
      * int foo(int a) {
-     *         return a; // when getType is invoked on "a" it returns the type "int"
-     *     }
+     * return a; // when getType is invoked on "a" it returns the type "int"
+     * }
      * </code>
-     *
+     * <p>
      * Now, users started using also of names of types itself, which do not have a type.
-     *
+     * <p>
      * For example:
      * <code>
      * class A {
-     *     int foo(int a) {
-     *         return A.someStaticField; // when getType is invoked on "A", which represents a class, it returns
-     *                                   // the type "A" itself while it used to throw UnsolvedSymbolException
-     *     }
+     * int foo(int a) {
+     * return A.someStaticField; // when getType is invoked on "A", which represents a class, it returns
+     * // the type "A" itself while it used to throw UnsolvedSymbolException
+     * }
      * </code>
-     *
+     * <p>
      * To accomodate this usage and avoid confusion this method return
      * the type itself when used on the name of type.
      */
@@ -284,7 +292,7 @@ public ResolvedType getType(Node node) {
             return getType(node, true);
         } catch (UnsolvedSymbolException e) {
             if (node instanceof NameExpr) {
-                NameExpr nameExpr = (NameExpr)node;
+                NameExpr nameExpr = (NameExpr) node;
                 SymbolReference<ResolvedTypeDeclaration> typeDeclaration = JavaParserFactory.getContext(node, typeSolver)
                         .solveType(nameExpr.getNameAsString(), typeSolver);
                 if (typeDeclaration.isSolved() && typeDeclaration.getCorrespondingDeclaration() instanceof ResolvedReferenceTypeDeclaration) {
@@ -298,66 +306,49 @@ public ResolvedType getType(Node node) {
 
     public ResolvedType getType(Node node, boolean solveLambdas) {
         if (solveLambdas) {
-            if (!cacheWithLambdasSolved.containsKey(node)) {
+            if (!node.containsData(TYPE_WITH_LAMBDAS_RESOLVED)) {
                 ResolvedType res = getTypeConcrete(node, solveLambdas);
 
-                cacheWithLambdasSolved.put(node, res);
+                node.setData(TYPE_WITH_LAMBDAS_RESOLVED, res);
 
                 boolean secondPassNecessary = false;
                 if (node instanceof MethodCallExpr) {
                     MethodCallExpr methodCallExpr = (MethodCallExpr) node;
                     for (Node arg : methodCallExpr.getArguments()) {
-                        if (!cacheWithLambdasSolved.containsKey(arg)) {
+                        if (!arg.containsData(TYPE_WITH_LAMBDAS_RESOLVED)) {
                             getType(arg, true);
                             secondPassNecessary = true;
                         }
                     }
                 }
                 if (secondPassNecessary) {
-                    cacheWithLambdasSolved.remove(node);
-                    cacheWithLambdasSolved.put(node, getType(node, true));
+                    node.removeData(TYPE_WITH_LAMBDAS_RESOLVED);
+                    ResolvedType type = getType(node, true);
+                    node.setData(TYPE_WITH_LAMBDAS_RESOLVED, type);
+
                 }
                 Log.trace("getType on %s  -> %s" ,node, res);
             }
-            return cacheWithLambdasSolved.get(node);
+            return node.getData(TYPE_WITH_LAMBDAS_RESOLVED);
         } else {
-            Optional<ResolvedType> res = find(cacheWithLambdasSolved, node);
+            Optional<ResolvedType> res = find(TYPE_WITH_LAMBDAS_RESOLVED, node);
             if (res.isPresent()) {
                 return res.get();
             }
-            res = find(cacheWithoutLambdasSolved, node);
+            res = find(TYPE_WITHOUT_LAMBDAS_RESOLVED, node);
             if (!res.isPresent()) {
                 ResolvedType resType = getTypeConcrete(node, solveLambdas);
-                cacheWithoutLambdasSolved.put(node, resType);
+                node.setData(TYPE_WITHOUT_LAMBDAS_RESOLVED, resType);
                 Log.trace("getType on %s (no solveLambdas) -> %s", node, res);
                 return resType;
             }
             return res.get();
         }
     }
 
-    private Optional<ResolvedType> find(Map<Node, ResolvedType> map, Node node) {
-        if (map.containsKey(node)) {
-            return Optional.of(map.get(node));
-        }
-        if (node instanceof LambdaExpr) {
-            return find(map, (LambdaExpr) node);
-        } else {
-            return Optional.empty();
-        }
-    }
-
-    /**
-     * For some reasons LambdaExprs are duplicate and the equals method is not implemented correctly.
-     */
-    private Optional<ResolvedType> find(Map<Node, ResolvedType> map, LambdaExpr lambdaExpr) {
-        for (Node key : map.keySet()) {
-            if (key instanceof LambdaExpr) {
-                LambdaExpr keyLambdaExpr = (LambdaExpr) key;
-                if (keyLambdaExpr.toString().equals(lambdaExpr.toString()) && requireParentNode(keyLambdaExpr) == requireParentNode(lambdaExpr)) {
-                    return Optional.of(map.get(keyLambdaExpr));
-                }
-            }
+    private Optional<ResolvedType> find(DataKey<ResolvedType> dataKey, Node node) {
+        if (node.containsData(dataKey)) {
+            return Optional.of(node.getData(dataKey));
         }
         return Optional.empty();
     }

javaparser-symbol-solver-core/src/main/java/com/github/javaparser/symbolsolver/javaparsermodel/JavaParserFactory.java

@@ -22,27 +22,17 @@
 import com.github.javaparser.ast.expr.*;
 import com.github.javaparser.ast.stmt.*;
 import com.github.javaparser.ast.type.TypeParameter;
-import com.github.javaparser.resolution.UnsolvedSymbolException;
 import com.github.javaparser.resolution.declarations.ResolvedReferenceTypeDeclaration;
-import com.github.javaparser.resolution.declarations.ResolvedValueDeclaration;
-import com.github.javaparser.resolution.types.ResolvedReferenceType;
-import com.github.javaparser.resolution.types.ResolvedType;
 import com.github.javaparser.symbolsolver.core.resolution.Context;
 import com.github.javaparser.symbolsolver.javaparsermodel.contexts.*;
-import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserAnnotationDeclaration;
-import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserClassDeclaration;
-import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserEnumDeclaration;
-import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserInterfaceDeclaration;
-import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserTypeParameter;
+import com.github.javaparser.symbolsolver.javaparsermodel.declarations.*;
 import com.github.javaparser.symbolsolver.javaparsermodel.declarators.FieldSymbolDeclarator;
 import com.github.javaparser.symbolsolver.javaparsermodel.declarators.NoSymbolDeclarator;
 import com.github.javaparser.symbolsolver.javaparsermodel.declarators.ParameterSymbolDeclarator;
 import com.github.javaparser.symbolsolver.javaparsermodel.declarators.VariableSymbolDeclarator;
-import com.github.javaparser.symbolsolver.model.resolution.SymbolReference;
 import com.github.javaparser.symbolsolver.model.resolution.TypeSolver;
 import com.github.javaparser.symbolsolver.resolution.SymbolDeclarator;
 
-import static com.github.javaparser.symbolsolver.javaparser.Navigator.getParentNode;
 import static com.github.javaparser.symbolsolver.javaparser.Navigator.requireParentNode;
 
 /**

javaparser-symbol-solver-testing/src/test/java/com/github/javaparser/symbolsolver/resolution/LambdaResolutionTest.java

@@ -16,7 +16,6 @@
 
 package com.github.javaparser.symbolsolver.resolution;
 
-import com.github.javaparser.ParseException;
 import com.github.javaparser.ast.CompilationUnit;
 import com.github.javaparser.ast.body.MethodDeclaration;
 import com.github.javaparser.ast.expr.Expression;
@@ -27,6 +26,7 @@
 import com.github.javaparser.symbolsolver.javaparser.Navigator;
 import com.github.javaparser.symbolsolver.javaparsermodel.JavaParserFacade;
 import com.github.javaparser.symbolsolver.resolution.typesolvers.ReflectionTypeSolver;
+import com.github.javaparser.utils.Log;
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;