identifier normalization

partiql · Jul 12, 2024 · 00412d8 · 00412d8
1 parent 4dc27b4
commit 00412d8
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 6 deletions.
diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/IdentifierManager.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/IdentifierManager.kt
@@ -0,0 +1,32 @@
+package org.partiql.planner.internal
+
+import org.partiql.ast.Binder
+import org.partiql.ast.Identifier
+import org.partiql.ast.binder
+import org.partiql.ast.identifierSymbol
+
+internal class IdentifierManager(
+    val casePreservation: Boolean,
+    val rValue: RValue
+) {
+    fun normalizeBinder(binder: Binder): Binder {
+        return if (casePreservation) {
+            binder(binder.symbol, Binder.CaseSensitivity.SENSITIVE)
+        }
+        // This is a conscious decision
+        // Even though the case preservation is off, we still need a way to store the binder
+        // To make this compatible with the current implementation,
+        // the decision here is: we normalize by preserving case.
+        else {
+            binder(binder.symbol, Binder.CaseSensitivity.SENSITIVE)
+        }
+    }
+
+    fun normalizeRvalue(id: Identifier.Symbol): Identifier.Symbol =
+        when (rValue) {
+            RValue.FOLDING_UP -> identifierSymbol(id.symbol.uppercase(), Identifier.CaseSensitivity.SENSITIVE)
+            RValue.FOLDING_DOWN -> identifierSymbol(id.symbol.lowercase(), Identifier.CaseSensitivity.SENSITIVE)
+            RValue.SENSITIVE -> identifierSymbol(id.symbol, Identifier.CaseSensitivity.SENSITIVE)
+            RValue.INSENSITIVE -> identifierSymbol(id.symbol, Identifier.CaseSensitivity.INSENSITIVE)
+        }
+}
diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PartiQLPlannerDefault.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PartiQLPlannerDefault.kt
@@ -27,7 +27,7 @@ internal class PartiQLPlannerDefault(
         val env = Env(session)
 
         // 1. Normalize
-        val ast = statement.normalize()
+        val ast = statement.normalize(flags)
 
         // 2. AST to Rel/Rex
         val root = AstToPlan.apply(ast, env)

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/astPasses/Normalize.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/astPasses/Normalize.kt
@@ -15,15 +15,21 @@
 package org.partiql.planner.internal.astPasses
 
 import org.partiql.ast.Statement
+import org.partiql.planner.internal.BooleanFlag
+import org.partiql.planner.internal.PlannerFlag
+import org.partiql.planner.internal.RValue
 
 /**
  * AST normalization
  */
-internal fun Statement.normalize(): Statement {
+internal fun Statement.normalize(flags: Set<PlannerFlag>): Statement {
     // could be a fold, but this is nice for setting breakpoints
     var ast = this
     ast = NormalizeFromSource.apply(ast)
     ast = NormalizeGroupBy.apply(ast)
     ast = NormalizeSelect.apply(ast)
+    val casePreservation = flags.any { it == BooleanFlag.CASE_PRESERVATION }
+    val rValue = flags.first { it is RValue } as RValue
+    ast = NormalizeIdentifier(casePreservation, rValue).apply(ast)
     return ast
 }
diff --git a/...iql-planner/src/main/kotlin/org/partiql/planner/internal/astPasses/NormalizeIdentifier.kt b/...iql-planner/src/main/kotlin/org/partiql/planner/internal/astPasses/NormalizeIdentifier.kt
@@ -1,5 +1,11 @@
 package org.partiql.planner.internal.astPasses
 
+import org.partiql.ast.Binder
+import org.partiql.ast.Identifier
+import org.partiql.ast.Statement
+import org.partiql.ast.normalize.AstPass
+import org.partiql.ast.util.AstRewriter
+import org.partiql.planner.internal.IdentifierManager
 import org.partiql.planner.internal.RValue
 
 /**
@@ -18,6 +24,18 @@ import org.partiql.planner.internal.RValue
  *   - Case Insensitive: Matching behavior, string comparison with case ignored.
  */
 internal class NormalizeIdentifier(
-    val casePreservation: Boolean,
-    val rValue: RValue
-)
+    casePreservation: Boolean,
+    rValue: RValue
+) : AstPass {
+
+    val identifierManager = IdentifierManager(casePreservation, rValue)
+    override fun apply(statement: Statement): Statement = Visitor(identifierManager).visitStatement(statement, Unit) as Statement
+
+    private class Visitor(val identifierManager: IdentifierManager) : AstRewriter<Unit>() {
+
+        override fun visitIdentifierSymbol(node: Identifier.Symbol, ctx: Unit) =
+            identifierManager.normalizeRvalue(node)
+
+        override fun visitBinder(node: Binder, ctx: Unit) = identifierManager.normalizeBinder(node)
+    }
+}
diff --git a/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTestsPorted.kt b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTestsPorted.kt
@@ -130,7 +130,12 @@ internal class PlanTyperTestsPorted {
     companion object {
 
         private val parser = PartiQLParser.default()
-        private val planner = PartiQLPlanner.builder().signalMode().build()
+        private val planner = PartiQLPlanner
+            .builder()
+            .signalMode()
+            .casePreserve()
+            .lookUpBehavior("INSENSITIVE")
+            .build()
 
         private fun assertProblemExists(problem: Problem) = ProblemHandler { problems, ignoreSourceLocation ->
             val message = buildString {