Fix compiler error for lambda parameter with non-letter

Fix a generated code error when lambda parameter contains a character that is not part of a valid Java identifier.
trinodb · Sep 9, 2020 · 8732f67 · 8732f67
1 parent cc289c9
commit 8732f67
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 2 deletions.
diff --git a/presto-main/src/main/java/io/prestosql/sql/gen/BytecodeUtils.java b/presto-main/src/main/java/io/prestosql/sql/gen/BytecodeUtils.java
@@ -420,14 +420,22 @@ else if (type == Boolean.class) {
     public static BytecodeExpression invoke(Binding binding, String name)
     {
         // ensure that name doesn't have a special characters
-        return invokeDynamic(BOOTSTRAP_METHOD, ImmutableList.of(binding.getBindingId()), name.replaceAll("[^A-Za-z0-9_$]", "_"), binding.getType());
+        return invokeDynamic(BOOTSTRAP_METHOD, ImmutableList.of(binding.getBindingId()), sanitizeName(name), binding.getType());
     }
 
     public static BytecodeExpression invoke(Binding binding, BoundSignature signature)
     {
         return invoke(binding, signature.getName());
     }
 
+    /**
+     * Replace characters that are not safe to use in a JVM identifier.
+     */
+    public static String sanitizeName(String name)
+    {
+        return name.replaceAll("[^A-Za-z0-9_$]", "_");
+    }
+
     public static BytecodeNode generateWrite(CallSiteBinder callSiteBinder, Scope scope, Variable wasNullVariable, Type type)
     {
         Class<?> valueJavaType = type.getJavaType();

diff --git a/presto-main/src/main/java/io/prestosql/sql/gen/LambdaBytecodeGenerator.java b/presto-main/src/main/java/io/prestosql/sql/gen/LambdaBytecodeGenerator.java
@@ -122,7 +122,7 @@ public static CompiledLambda preGenerateLambdaExpression(
         for (int i = 0; i < lambdaExpression.getArguments().size(); i++) {
             Class<?> type = Primitives.wrap(lambdaExpression.getArgumentTypes().get(i).getJavaType());
             String argumentName = lambdaExpression.getArguments().get(i);
-            Parameter arg = arg("lambda_" + argumentName, type);
+            Parameter arg = arg("lambda_" + i + "_" + BytecodeUtils.sanitizeName(argumentName), type);
             parameters.add(arg);
             parameterMapBuilder.put(argumentName, new ParameterAndType(arg, type));
         }

diff --git a/presto-main/src/test/java/io/prestosql/operator/scalar/TestLambdaExpression.java b/presto-main/src/test/java/io/prestosql/operator/scalar/TestLambdaExpression.java
@@ -54,6 +54,14 @@ public void testBasic()
         assertFunction("apply(5 + RANDOM(1), x -> x + 1)", INTEGER, 6);
     }
 
+    @Test
+    public void testParameterName()
+    {
+        // parameter which is not valid identifier in Java
+        String nonLetters = "a.b c; d ' \n \\n \"";
+        assertFunction("apply(5, " + quote(nonLetters) + " -> " + quote(nonLetters) + " * 2)", INTEGER, 10);
+    }
+
     @Test
     public void testNull()
     {
@@ -158,4 +166,9 @@ public void testTypeCombinations()
         assertFunction("apply(ARRAY['abc', NULL, '123'], x -> x[2])", createVarcharType(3), null);
         assertFunction("apply(MAP(ARRAY['abc', 'def'], ARRAY[123, 456]), x -> map_keys(x))", new ArrayType(createVarcharType(3)), ImmutableList.of("abc", "def"));
     }
+
+    private static String quote(String identifier)
+    {
+        return "\"" + identifier.replace("\"", "\"\"") + "\"";
+    }
 }
diff --git a/presto-main/src/test/java/io/prestosql/sql/query/TestLambdaExpressions.java b/presto-main/src/test/java/io/prestosql/sql/query/TestLambdaExpressions.java
@@ -63,6 +63,20 @@ public void testDuplicateLambdaExpressions()
                 .matches("VALUES ROW(ARRAY[2, 3, 4], ARRAY[11e0, 21e0, 31e0])");
     }
 
+    @Test
+    public void testParameterName()
+    {
+        // lambda may be using parameter which is not valid identifier in Java
+        assertThat(assertions.query("" +
+                "WITH t AS ( " +
+                "    SELECT count(*) AS \"a.b c; d\" FROM (VALUES (42)) " +
+                "    UNION ALL " +
+                "    SELECT * FROM (VALUES (77)) v(\"a.b c; d\") " +
+                ") " +
+                "SELECT transform(ARRAY[1], x -> x + \"a.b c; d\") FROM t"))
+                .matches("VALUES ARRAY[BIGINT '2'], ARRAY[BIGINT '78']");
+    }
+
     @Test
     public void testNestedLambda()
     {