Add nodes to handle types

modules/lang-painless/src/main/java/org/elasticsearch/painless/antlr/Walker.java

@@ -109,6 +109,7 @@
 import org.elasticsearch.painless.node.AExpression;
 import org.elasticsearch.painless.node.ANode;
 import org.elasticsearch.painless.node.AStatement;
+import org.elasticsearch.painless.node.DUnresolvedType;
 import org.elasticsearch.painless.node.EAssignment;
 import org.elasticsearch.painless.node.EBinary;
 import org.elasticsearch.painless.node.EBool;
@@ -478,8 +479,9 @@ public ANode visitDeclaration(DeclarationContext ctx) {
         for (DeclvarContext declvar : ctx.declvar()) {
             String name = declvar.ID().getText();
             AExpression expression = declvar.expression() == null ? null : (AExpression)visit(declvar.expression());
+            DUnresolvedType unresolvedType = new DUnresolvedType(location(declvar), type);
 
-            declarations.add(new SDeclaration(location(declvar), type, name, expression));
+            declarations.add(new SDeclaration(location(declvar), unresolvedType, name, expression));
         }
 
         return new SDeclBlock(location(ctx), declarations);

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/DResolvedType.java

@@ -0,0 +1,81 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.node;
+
+import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.lookup.PainlessLookup;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
+
+import java.util.Objects;
+
+/**
+ * Represents a Painless type as a {@link Class}. This may still
+ * require resolution to ensure the type in the {@link PainlessLookup}.
+ */
+public class DResolvedType extends DType {
+
+    protected final Class<?> type;
+
+    /**
+     * If set to {@code true} ensures the type is in the {@link PainlessLookup}.
+     * If set to {@code false} assumes the type is valid.
+     */
+    protected final boolean checkInLookup;
+
+    public DResolvedType(Location location, Class<?> type) {
+        this(location, type, true);
+    }
+
+    public DResolvedType(Location location, Class<?> type, boolean checkInLookup) {
+        super(location);
+        this.type = Objects.requireNonNull(type);
+        this.checkInLookup = checkInLookup;
+    }
+
+    /**
+     * If {@link #checkInLookup} is {@code true} checks if the type is in the
+     * {@link PainlessLookup}, otherwise returns {@code this}.
+     * @throws IllegalArgumentException if both checking the type is in the {@link PainlessLookup}
+     * and the type cannot be resolved from the {@link PainlessLookup}
+     * @return a {@link DResolvedType} where the resolved Painless type is retrievable
+     */
+    @Override
+    public DResolvedType resolveType(PainlessLookup painlessLookup) {
+        if (checkInLookup == false) {
+            return this;
+        }
+
+        if (painlessLookup.getClasses().contains(type) == false) {
+            throw location.createError(new IllegalArgumentException(
+                    "cannot resolve type [" + PainlessLookupUtility.typeToCanonicalTypeName(type) + "]"));
+        }
+
+        return new DResolvedType(location, type, false);
+    }
+
+    public Class<?> getType() {
+        return type;
+    }
+
+    @Override
+    public String toString() {
+        return " (DResolvedType [" + PainlessLookupUtility.typeToCanonicalTypeName(type) + "])";
+    }
+}

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/DType.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.node;
+
+import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.lookup.PainlessLookup;
+
+import java.util.Objects;
+
+/**
+ * Represents an abstract Painless type. {@link DType} nodes must be
+ * resolved using {@link #resolveType(PainlessLookup)} to retrieve the
+ * actual Painless type. {@link DType} exists as a base class so consumers
+ * may have either a {@link DUnresolvedType} representing a Painless
+ * canonical type name or a {@link DResolvedType} representing a Painless
+ * type as the Painless AST is constructed. This allows Painless types already
+ * resolved at the time of Painless AST construction to not be forced to
+ * convert back to a Painless canonical type name and then re-resolved.
+ */
+public abstract class DType {
+
+    protected final Location location;
+
+    public DType(Location location) {
+        this.location = Objects.requireNonNull(location);
+    }
+
+    public abstract DResolvedType resolveType(PainlessLookup painlessLookup);
+}

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/DUnresolvedType.java

@@ -0,0 +1,61 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.node;
+
+import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.lookup.PainlessLookup;
+
+import java.util.Objects;
+
+/**
+ * Represents a canonical Painless type name as a {@link String}
+ * that requires resolution.
+ */
+public class DUnresolvedType extends DType {
+
+    protected final String typeName;
+
+    public DUnresolvedType(Location location, String typeName) {
+        super(location);
+        this.typeName = Objects.requireNonNull(typeName);
+    }
+
+    /**
+     * Resolves the canonical Painless type name to a Painless type.
+     * @throws IllegalArgumentException if the type cannot be resolved from the {@link PainlessLookup}
+     * @return a {@link DResolvedType} where the resolved Painless type is retrievable
+     */
+    @Override
+    public DResolvedType resolveType(PainlessLookup painlessLookup) {
+        Class<?> type = painlessLookup.canonicalTypeNameToType(typeName);
+
+        if (type == null) {
+            throw location.createError(new IllegalArgumentException("cannot resolve type [" + typeName + "]"));
+        }
+
+        return new DResolvedType(location, type);
+    }
+
+    @Override
+    public String toString() {
+        return "(DUnresolvedType [" + typeName + "])";
+    }
+}
+

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SDeclaration.java

@@ -36,13 +36,13 @@
  */
 public final class SDeclaration extends AStatement {
 
-    private final String type;
+    private final DType type;
     private final String name;
     private AExpression expression;
 
     private Variable variable = null;
 
-    public SDeclaration(Location location, String type, String name, AExpression expression) {
+    public SDeclaration(Location location, DType type, String name, AExpression expression) {
         super(location);
 
         this.type = Objects.requireNonNull(type);
@@ -61,19 +61,15 @@ void extractVariables(Set<String> variables) {
 
     @Override
     void analyze(ScriptRoot scriptRoot, Locals locals) {
-        Class<?> clazz = scriptRoot.getPainlessLookup().canonicalTypeNameToType(this.type);
-
-        if (clazz == null) {
-            throw createError(new IllegalArgumentException("Not a type [" + this.type + "]."));
-        }
+        DResolvedType resolvedType = type.resolveType(scriptRoot.getPainlessLookup());
 
         if (expression != null) {
-            expression.expected = clazz;
+            expression.expected = resolvedType.getType();
             expression.analyze(scriptRoot, locals);
             expression = expression.cast(scriptRoot, locals);
         }
 
-        variable = locals.addVariable(location, clazz, name, false);
+        variable = locals.addVariable(location, resolvedType.getType(), name, false);
     }
 
     @Override