From f89150ed69af799aefdea96122784789963e60cc Mon Sep 17 00:00:00 2001 From: Stuart Tettemer Date: Wed, 24 Jan 2024 16:44:41 -0600 Subject: [PATCH] start tracking variables --- .../org/elasticsearch/painless/Compiler.java | 10 +-- .../phase/DefaultTypeAnalysisPhase.java | 15 +++- .../painless/symbol/Variable.java | 5 ++ .../painless/symbol/VariableScope.java | 85 +++++++++++++------ .../painless/TypeInferenceTests.java | 12 +++ 5 files changed, 95 insertions(+), 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 index 22fbea8c28701..a79465e858c98 100644 --- 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/modules/lang-painless/src/main/java/org/elasticsearch/painless/phase/DefaultTypeAnalysisPhase.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/phase/DefaultTypeAnalysisPhase.java index 6f96021fae2a8..1f81cd2f95412 100644 --- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/phase/DefaultTypeAnalysisPhase.java +++ b/modules/lang-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 index 0e9005cf8cccb..b055b6cd744c2 100644 --- 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 index fa564091a8b39..f71057343709e 100644 --- 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 variables = new HashMap<>(); // assignments to key from typed expressions - Map> typedExprAssignments = new HashMap<>(); + Map> typedExprAssignments = new HashMap<>(); + Map> typedVarAssignments = new HashMap<>(); // assignments to key from def expressions Map> defExprAssignment = new HashMap<>(); @@ -42,12 +44,26 @@ public abstract class VariableScope { Map> 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> 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> 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 index 540e5a805eb27..0b7bf5e2a4493 100644 --- 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"); + } }