Pull out semantic analysis phase

stu-elastic · Jan 24, 2024 · 7119c92 · 7119c92
1 parent bc065cd
commit 7119c92
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 24 deletions.
diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/Compiler.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/Compiler.java
@@ -211,13 +211,9 @@ private static void addFactoryMethod(Map<String, Class<?>> additionalClasses, Cl
      */
     ScriptScope compile(Loader loader, String name, String source, CompilerSettings settings) {
         String scriptName = Location.computeSourceName(name);
-        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SClass root = Walker.buildPainlessTree(scriptName, source, settings);
-        ScriptScope scriptScope = new ScriptScope(painlessLookup, settings, scriptClassInfo, scriptName, source, root.getIdentifier() + 1);
-        new PainlessSemanticHeaderPhase().visitClass(root, scriptScope);
-        VariableScope variableScope = VariableScope.programScope(root.getLocation(), scriptScope);
-        new DefaultTypeAnalysisPhase().visitClass(root, variableScope);
-        new PainlessSemanticAnalysisPhase().visitClass(root, scriptScope);
+        ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
+        ScriptScope typeScriptScope = runTypeAnalysis(scriptName, root, source, settings);
         new PainlessUserTreeToIRTreePhase().visitClass(root, scriptScope);
         ClassNode classNode = (ClassNode) scriptScope.getDecoration(root, IRNodeDecoration.class).irNode();
         new DefaultStringConcatenationOptimizationPhase().visitClass(classNode, null);
@@ -249,13 +245,9 @@ ScriptScope compile(Loader loader, String name, String source, CompilerSettings
      */
     byte[] compile(String name, String source, CompilerSettings settings, Printer debugStream) {
         String scriptName = Location.computeSourceName(name);
-        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SClass root = Walker.buildPainlessTree(scriptName, source, settings);
-        ScriptScope scriptScope = new ScriptScope(painlessLookup, settings, scriptClassInfo, scriptName, source, root.getIdentifier() + 1);
-        new PainlessSemanticHeaderPhase().visitClass(root, scriptScope);
-        VariableScope variableScope = VariableScope.programScope(root.getLocation(), scriptScope);
-        new DefaultTypeAnalysisPhase().visitClass(root, variableScope);
-        new PainlessSemanticAnalysisPhase().visitClass(root, scriptScope);
+        ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
+        ScriptScope typeScriptScope = runTypeAnalysis(scriptName, root, source, settings);
         new PainlessUserTreeToIRTreePhase().visitClass(root, scriptScope);
         ClassNode classNode = (ClassNode) scriptScope.getDecoration(root, IRNodeDecoration.class).irNode();
         new DefaultStringConcatenationOptimizationPhase().visitClass(classNode, null);
@@ -281,14 +273,9 @@ byte[] compile(
         IRTreeVisitor<WriteScope> asmPhaseVisitor
     ) {
         String scriptName = Location.computeSourceName(name);
-        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SClass root = Walker.buildPainlessTree(scriptName, source, settings);
-        ScriptScope scriptScope = new ScriptScope(painlessLookup, settings, scriptClassInfo, scriptName, source, root.getIdentifier() + 1);
-
-        new PainlessSemanticHeaderPhase().visitClass(root, scriptScope);
-        VariableScope variableScope = VariableScope.programScope(root.getLocation(), scriptScope);
-        new DefaultTypeAnalysisPhase().visitClass(root, variableScope);
-        new PainlessSemanticAnalysisPhase().visitClass(root, scriptScope);
+        ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
+        ScriptScope typeScriptScope = runTypeAnalysis(scriptName, root, source, settings);
         if (semanticPhaseVisitor != null) {
             semanticPhaseVisitor.visitClass(root, scriptScope);
         }
@@ -313,4 +300,19 @@ byte[] compile(
 
         return classNode.getBytes();
     }
+
+    ScriptScope runTypeAnalysis(String scriptName, SClass root, String source, CompilerSettings settings) {
+        ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
+        VariableScope variableScope = VariableScope.programScope(root.getLocation(), scriptScope);
+        new DefaultTypeAnalysisPhase().visitClass(root, variableScope);
+        return scriptScope;
+    }
+
+    ScriptScope runSemanticAnalysisPhase(String scriptName, SClass root, String source, CompilerSettings settings) {
+        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
+        ScriptScope scriptScope = new ScriptScope(painlessLookup, settings, scriptClassInfo, scriptName, source, root.getIdentifier() + 1);
+        new PainlessSemanticHeaderPhase().visitClass(root, scriptScope);
+        new PainlessSemanticAnalysisPhase().visitClass(root, scriptScope);
+        return scriptScope;
+    }
 }
diff --git a/...ang-painless/src/main/java/org/elasticsearch/painless/phase/DefaultTypeAnalysisPhase.java b/...ang-painless/src/main/java/org/elasticsearch/painless/phase/DefaultTypeAnalysisPhase.java
@@ -55,20 +55,29 @@ public void visitFunction(SFunction userFunctionNode, VariableScope variableScop
     @Override
     public void visitAssignment(EAssignment userAssignmentNode, VariableScope variableScope) {
         AExpression userLeftNode = userAssignmentNode.getLeftNode();
-        //Class<?> leftValueType = variableScope.getValueType(userLeftNode);
+        Class<?> leftValueType = variableScope.getValueType(userLeftNode);
         visit(userLeftNode, variableScope);
 
         AExpression userRightNode = userAssignmentNode.getRightNode();
         visit(userRightNode, variableScope);
-        //Class<?> rightValueType = variableScope.getValueType(userRightNode);
+        Class<?> rightValueType = variableScope.getValueType(userRightNode);
 
         //Operation operation = userAssignmentNode.getOperation();
 
     }
 
     @Override
     public void visitDeclaration(SDeclaration userDeclarationNode, VariableScope variableScope) {
-        super.visitDeclaration(userDeclarationNode, variableScope);
+        String symbol = userDeclarationNode.getSymbol();
+        String canonicalTypeName = userDeclarationNode.getCanonicalTypeName();
+        Class<?> type = variableScope.canonicalTypeNameToType(canonicalTypeName);
+        AExpression userValueNode = userDeclarationNode.getValueNode();
+        Class<?> valueType = null;
+        if (userValueNode != null) {
+            visit(userValueNode, variableScope);
+            valueType = variableScope.getValueType(userValueNode);
+        }
+        variableScope.defineVariable(userDeclarationNode.getLocation(), type, symbol, false);
     }
 
     @Override

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/Decorator.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/Decorator.java
@@ -24,8 +24,8 @@ public interface Condition {
 
     }
 
-    private final ArrayList<Map<Class<? extends Decoration>, Decoration>> decorations;
-    private final ArrayList<Set<Class<? extends Condition>>> conditions;
+    protected final ArrayList<Map<Class<? extends Decoration>, Decoration>> decorations;
+    protected final ArrayList<Set<Class<? extends Condition>>> conditions;
 
     public Decorator(int nodeCount) {
         decorations = new ArrayList<>(nodeCount);

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/VariableScope.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/VariableScope.java
@@ -89,6 +89,10 @@ public Class<?> getValueType(ANode node) {
         return parent.getValueType(node);
     }
 
+    public Class<?> canonicalTypeNameToType(String canonicalTypeName) {
+        return parent.canonicalTypeNameToType(canonicalTypeName);
+    }
+
     public void addVarAssignment(String target, String source, Location location) {
         ScopedVariable targetScope = getScopedVariable(target, location);
         ScopedVariable sourceScope = getScopedVariable(source, location);
@@ -161,6 +165,16 @@ public List<Class<?>> getFunctionTypeParameters(String name, int arity) {
             return function.getTypeParameters();
         }
 
+        @Override
+        public Class<?> canonicalTypeNameToType(String canonicalTypeName) {
+            return scriptScope.getPainlessLookup().canonicalTypeNameToType(canonicalTypeName);
+        }
+
+        @Override
+        public Class<?> getValueType(ANode node) {
+            return scriptScope.getDecoration(node, Decorations.ValueType.class).valueType();
+        }
+
         @Override
         public VariableScope functionScope(Location location) {
             return addScope(new FunctionScope(location, this));

diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/TypeInferenceTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/TypeInferenceTests.java
@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.painless;
+
+public class TypeInferenceTests extends ScriptTestCase {
+    public void testAssignmentTypeInference() {
+        var script = """
+            def foo = 1;
+            if (foo > 0) {
+                return "greater";
+            }
+            return "lessthan";
+            """;
+        assertBytecodeExists(script, "IF_ICMPGT L2");
+    }
+}