start tracking variables

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

Hide file tree

Showing 5 changed files with 95 additions and 32 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
@@ -213,7 +213,7 @@ ScriptScope compile(Loader loader, String name, String source, CompilerSettings
         String scriptName = Location.computeSourceName(name);
         SClass root = Walker.buildPainlessTree(scriptName, source, settings);
         ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
-        ScriptScope typeScriptScope = runTypeAnalysis(scriptName, root, source, settings);
+        VariableScope variableScope = runTypeAnalysis(scriptName, root, source, settings);
         new PainlessUserTreeToIRTreePhase().visitClass(root, scriptScope);
         ClassNode classNode = (ClassNode) scriptScope.getDecoration(root, IRNodeDecoration.class).irNode();
         new DefaultStringConcatenationOptimizationPhase().visitClass(classNode, null);
@@ -247,7 +247,7 @@ byte[] compile(String name, String source, CompilerSettings settings, Printer de
         String scriptName = Location.computeSourceName(name);
         SClass root = Walker.buildPainlessTree(scriptName, source, settings);
         ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
-        ScriptScope typeScriptScope = runTypeAnalysis(scriptName, root, source, settings);
+        VariableScope variableScope = runTypeAnalysis(scriptName, root, source, settings);
         new PainlessUserTreeToIRTreePhase().visitClass(root, scriptScope);
         ClassNode classNode = (ClassNode) scriptScope.getDecoration(root, IRNodeDecoration.class).irNode();
         new DefaultStringConcatenationOptimizationPhase().visitClass(classNode, null);
@@ -275,7 +275,7 @@ byte[] compile(
         String scriptName = Location.computeSourceName(name);
         SClass root = Walker.buildPainlessTree(scriptName, source, settings);
         ScriptScope scriptScope = runSemanticAnalysisPhase(scriptName, root, source, settings);
-        ScriptScope typeScriptScope = runTypeAnalysis(scriptName, root, source, settings);
+        VariableScope variableScope = runTypeAnalysis(scriptName, root, source, settings);
         if (semanticPhaseVisitor != null) {
             semanticPhaseVisitor.visitClass(root, scriptScope);
         }
@@ -301,11 +301,11 @@ byte[] compile(
         return classNode.getBytes();
     }
 
-    ScriptScope runTypeAnalysis(String scriptName, SClass root, String source, CompilerSettings settings) {
+    VariableScope 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;
+        return variableScope;
     }
 
     ScriptScope runSemanticAnalysisPhase(String scriptName, SClass root, String source, CompilerSettings settings) {

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
@@ -8,10 +8,11 @@
 
 package org.elasticsearch.painless.phase;
 
-import org.elasticsearch.painless.Operation;
+import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.node.AExpression;
 import org.elasticsearch.painless.node.ANode;
 import org.elasticsearch.painless.node.EAssignment;
+import org.elasticsearch.painless.node.ESymbol;
 import org.elasticsearch.painless.node.SBlock;
 import org.elasticsearch.painless.node.SCatch;
 import org.elasticsearch.painless.node.SClass;
@@ -24,6 +25,7 @@
 import org.elasticsearch.painless.node.SIfElse;
 import org.elasticsearch.painless.node.STry;
 import org.elasticsearch.painless.node.SWhile;
+import org.elasticsearch.painless.symbol.Variable;
 import org.elasticsearch.painless.symbol.VariableScope;
 
 import java.util.List;
@@ -77,7 +79,16 @@ public void visitDeclaration(SDeclaration userDeclarationNode, VariableScope var
             visit(userValueNode, variableScope);
             valueType = variableScope.getValueType(userValueNode);
         }
-        variableScope.defineVariable(userDeclarationNode.getLocation(), type, symbol, false);
+        Location location = userDeclarationNode.getLocation();
+        Variable var = variableScope.defineVariable(location, type, symbol, false);
+        if (valueType == null) {
+            return;
+        }
+        if (userValueNode instanceof ESymbol source) {
+            variableScope.addVarAssignment(symbol, source.getSymbol(), location);
+        } else {
+            variableScope.addExprAssignment(symbol, userValueNode, location);
+        }
     }
 
     @Override

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/Variable.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/symbol/Variable.java
@@ -10,6 +10,7 @@
 
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
+import org.elasticsearch.painless.lookup.def;
 
 import java.util.Objects;
 
@@ -32,4 +33,8 @@ public record Variable(Class<?> type, String name, boolean isFinal, Location loc
     public String getCanonicalTypeName() {
         return PainlessLookupUtility.typeToCanonicalTypeName(type);
     }
+
+    public boolean isDef() {
+        return def.class == type;
+    }
 }
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
@@ -19,6 +19,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 
 public abstract class VariableScope {
@@ -30,7 +31,8 @@ public abstract class VariableScope {
     Map<String, Variable> variables = new HashMap<>();
 
     // assignments to key from typed expressions
-    Map<String, List<TypedAssignment>> typedExprAssignments = new HashMap<>();
+    Map<String, List<TypedExprAssignment>> typedExprAssignments = new HashMap<>();
+    Map<String, List<TypedVarAssignment>> typedVarAssignments = new HashMap<>();
 
     // assignments to key from def expressions
     Map<String, List<DefExprAssignment>> defExprAssignment = new HashMap<>();
@@ -42,12 +44,26 @@ public abstract class VariableScope {
     Map<String, List<DefVarAssignment>> defVarDependents = new HashMap<>();
     // in this case key is DefVarAssignment.source
 
+    // Assignment from a typed expression
+    private record TypedExprAssignment(Variable target, AExpression expr, Class<?> type, Location location) {
+        TypedExprAssignment {
+            Objects.requireNonNull(type);
+        }
+    }
+
+    // Assignment from a typed variable
+    private record TypedVarAssignment(Variable target, Variable source, Class<?> type, Location location) {
+        TypedVarAssignment {
+            Objects.requireNonNull(type);
+            if (type == def.class) {
+                throw new IllegalArgumentException("TypedVarAssignment for [" + target + "] cannot be def type");
+            }
+        }
+    }
+
     // Assignment from a def variable
     private record DefVarAssignment(Variable target, Variable source, Location location) {}
 
-    // Assignment from a typed expression
-    private record TypedAssignment(Variable target, AExpression expr, Location location) {}
-
     // Assignment from a def expression
     private record DefExprAssignment(Variable target, AExpression expr, Location location) {}
 
@@ -81,16 +97,17 @@ protected void addDefVarAssignment(DefVarAssignment assignment) {
         defVarAssignment.computeIfAbsent(assignment.target.name(), k -> new ArrayList<>()).add(assignment);
     }
 
-    public List<Class<?>> getFunctionTypeParameters(String name, int arity) {
-        return parent.getFunctionTypeParameters(name, arity);
-    }
-
-    public Class<?> getValueType(ANode node) {
-        return parent.getValueType(node);
-    }
-
-    public Class<?> canonicalTypeNameToType(String canonicalTypeName) {
-        return parent.canonicalTypeNameToType(canonicalTypeName);
+    public void addExprAssignment(String target, AExpression expr, Location location) {
+        ScopedVariable targetScope = getScopedVariable(target, location);
+        Class<?> valueType = getValueType(expr);
+        if (valueType == null) {
+            return; // TODO(stu): is this an error?
+        }
+        if (valueType == def.class) {
+            targetScope.scope.addDefExprAssignment(targetScope.variable, expr, location);
+        } else {
+            targetScope.scope.addTypedExprAssignment(targetScope.variable, expr, valueType, location);
+        }
     }
 
     public void addVarAssignment(String target, String source, Location location) {
@@ -101,20 +118,38 @@ public void addVarAssignment(String target, String source, Location location) {
             targetScope.scope.addDefVarAssignment(assignment);
             sourceScope.scope.addDefVarDependent(assignment);
         } else {
-            // TODO(stu) is this an typed expr assignment?
-            int a = 1;
-            if (a < 2) {
-                a++;
-            }
+            targetScope.scope.addTypedVarAssignment(targetScope.variable, sourceScope.variable, sourceScope.variable.type(), location);
         }
     }
 
-    public void addExprAssignment(String target, AExpression expr, Location location) {
-        // TODO(stu)
-        int a = 1;
-        if (a < 2) {
-            a++;
-        }
+    protected void addDefExprAssignment(Variable variable, AExpression expr, Location location) {
+        defExprAssignment.computeIfAbsent(variable.name(), k -> new ArrayList<>()).add(
+            new DefExprAssignment(variable, expr, location)
+        );
+    }
+
+    protected void addTypedExprAssignment(Variable variable, AExpression expr, Class<?> type, Location location) {
+        typedExprAssignments.computeIfAbsent(variable.name(), k -> new ArrayList<>()).add(
+            new TypedExprAssignment(variable, expr, type, location)
+        );
+    }
+
+    protected void addTypedVarAssignment(Variable target, Variable source, Class<?> type, Location location) {
+        typedVarAssignments.computeIfAbsent(target.name(), k -> new ArrayList<>()).add(
+            new TypedVarAssignment(target, source, type, location)
+        );
+    }
+
+    public List<Class<?>> getFunctionTypeParameters(String name, int arity) {
+        return parent.getFunctionTypeParameters(name, arity);
+    }
+
+    public Class<?> getValueType(ANode node) {
+        return parent.getValueType(node);
+    }
+
+    public Class<?> canonicalTypeNameToType(String canonicalTypeName) {
+        return parent.canonicalTypeNameToType(canonicalTypeName);
     }
 
     public Variable defineVariable(Location location, Class<?> type, String name, boolean isReadOnly) {

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
@@ -19,4 +19,16 @@ public void testAssignmentTypeInference() {
             """;
         assertBytecodeExists(script, "IF_ICMPGT L2");
     }
+
+    public void testAssignmentTypeInferenceIndirect() {
+        var script = """
+            def bar = 1;
+            def foo = bar;
+            if (foo > 0) {
+                return "greater";
+            }
+            return "lessthan";
+            """;
+        assertBytecodeExists(script, "IF_ICMPGT L2");
+    }
 }