Skip to main content

pgevolve_core/parse/
normalize_expr.rs

1//! Build [`NormalizedExpr`] values from `pg_query` AST nodes.
2//!
3//! Phase-2 scope (per `docs/superpowers/plans/phase-2-parser.md`):
4//!
5//! - Strip redundant casts to a column's own type. `42::integer` for an integer
6//!   column collapses to `42`; `'foo'::text` for a text column collapses to `'foo'`.
7//! - Lowercase reserved keywords on the canonical text emitted by `pg_query`'s deparser.
8//! - Compute the BLAKE3 hash of the canonical text.
9//!
10//! Deferred to follow-up phase-2 issues (only affect equivalence sensitivity, not
11//! correctness — equivalent inputs that exercise these may diff for now):
12//!
13//! - Paren folding (collapsing trivial nested `A_Expr` parens).
14//! - Sorting commutative operands of `+`, `*`, `AND`, `OR`.
15
16use pg_query::NodeEnum;
17use pg_query::protobuf::{self, Node, ResTarget};
18
19use crate::ir::column_type::ColumnType;
20use crate::ir::default_expr::NormalizedExpr;
21use crate::parse::error::{ParseError, SourceLocation};
22
23/// Build a [`NormalizedExpr`] from a `pg_query` expression node.
24///
25/// `target_type`, when supplied, enables redundant-cast stripping: a `TypeCast`
26/// to that type is replaced by its inner expression.
27pub fn from_pg_node(
28    node: &NodeEnum,
29    target_type: Option<&ColumnType>,
30    location: &SourceLocation,
31) -> Result<NormalizedExpr, ParseError> {
32    let normalized = strip_redundant_cast(node.clone(), target_type);
33    let raw = deparse_expr(&normalized).map_err(|e| ParseError::Structural {
34        location: location.clone(),
35        message: format!("could not deparse expression: {e}"),
36    })?;
37    let canonical_text = lowercase_keywords(&raw);
38    Ok(NormalizedExpr::from_text(canonical_text))
39}
40
41/// If `node` is a `TypeCast` whose target type matches `target_type`, return
42/// the inner argument; otherwise return `node` unchanged.
43fn strip_redundant_cast(node: NodeEnum, target_type: Option<&ColumnType>) -> NodeEnum {
44    let Some(target) = target_type else {
45        return node;
46    };
47    let NodeEnum::TypeCast(cast) = &node else {
48        return node;
49    };
50    let Some(type_name) = cast.type_name.as_ref() else {
51        return node;
52    };
53    let Some(type_str) = render_type_name(type_name) else {
54        return node;
55    };
56    let Ok(parsed) = ColumnType::parse_from_pg_type_string(&type_str) else {
57        return node;
58    };
59    if &parsed != target {
60        return node;
61    }
62    // Strip: replace the cast with its inner argument.
63    cast.arg
64        .as_ref()
65        .and_then(|inner| inner.node.as_ref())
66        .cloned()
67        .unwrap_or(node)
68}
69
70/// Walk a `TypeName`'s `names` list and join the string fragments with dots.
71/// Returns `None` if any element is not a `String` node.
72fn render_type_name(type_name: &protobuf::TypeName) -> Option<String> {
73    let mut parts = Vec::with_capacity(type_name.names.len());
74    for n in &type_name.names {
75        let Some(NodeEnum::String(s)) = &n.node else {
76            return None;
77        };
78        parts.push(s.sval.clone());
79    }
80    if parts.is_empty() {
81        return None;
82    }
83    // pg_query stores types like `pg_catalog.int4`; we only care about the last
84    // segment for matching against [`ColumnType::parse_from_pg_type_string`],
85    // since that already understands aliases like `int4` → `Integer`.
86    Some(parts.last().cloned().unwrap_or_default())
87}
88
89/// Wrap an expression in `SELECT <expr>` and deparse, then strip the `SELECT `
90/// prefix. This is the workaround for `pg_query`'s deparse expecting a top-level
91/// statement node — there is no public `deparse_expr` entry point in v6.
92///
93/// The `protobuf::ParseResult.version` field must match `libpg_query`'s
94/// embedded `PG_VERSION_NUM`, otherwise the C deparser asserts and aborts the
95/// process. We borrow that version from a freshly-parsed `SELECT 1` rather than
96/// hard-coding it.
97fn deparse_expr(node: &NodeEnum) -> Result<String, pg_query::Error> {
98    let mut scaffold = pg_query::parse("SELECT 1")?.protobuf;
99    let raw = scaffold
100        .stmts
101        .get_mut(0)
102        .ok_or_else(|| pg_query::Error::Parse("scaffold has no stmts".into()))?;
103    let select_node = raw
104        .stmt
105        .as_mut()
106        .ok_or_else(|| pg_query::Error::Parse("scaffold stmt is None".into()))?
107        .node
108        .as_mut()
109        .ok_or_else(|| pg_query::Error::Parse("scaffold node is None".into()))?;
110    let NodeEnum::SelectStmt(select) = select_node else {
111        return Err(pg_query::Error::Parse(
112            "scaffold parse did not yield SelectStmt".into(),
113        ));
114    };
115    select.target_list = vec![Node {
116        node: Some(NodeEnum::ResTarget(Box::new(ResTarget {
117            name: String::new(),
118            indirection: vec![],
119            val: Some(Box::new(Node {
120                node: Some(node.clone()),
121            })),
122            location: -1,
123        }))),
124    }];
125    let s = pg_query::deparse(&scaffold)?;
126    Ok(s.trim_start_matches("SELECT ").to_string())
127}
128
129/// Reserved Postgres keywords that should appear lowercased in canonical text.
130/// We deliberately keep this small — `pg_query`'s deparser already emits most
131/// keywords lowercased; this list is a belt-and-suspenders pass for any node
132/// kinds where the deparser preserves the source's casing.
133const RESERVED_FUNC_KEYWORDS: &[&str] = &[
134    "AND", "OR", "NOT", "NULL", "TRUE", "FALSE", "IS", "IN", "LIKE", "BETWEEN", "CASE", "WHEN",
135    "THEN", "ELSE", "END", "CAST", "AS", "DISTINCT", "FROM", "WHERE", "ORDER", "BY", "GROUP",
136    "HAVING", "LIMIT", "OFFSET", "ASC", "DESC", "NULLS", "FIRST", "LAST", "WITH", "USING",
137    "COLLATE",
138];
139
140/// Lowercase whole-word reserved keywords in `s`. Quoted-string contents are not
141/// modified.
142fn lowercase_keywords(s: &str) -> String {
143    let mut out = String::with_capacity(s.len());
144    let mut chars = s.chars().peekable();
145    while let Some(c) = chars.next() {
146        if c == '\'' {
147            // Pass through quoted string bodies verbatim, including doubled `''`.
148            out.push(c);
149            while let Some(&n) = chars.peek() {
150                chars.next();
151                out.push(n);
152                if n == '\'' {
153                    if chars.peek() == Some(&'\'') {
154                        // SAFETY: peek() returned Some, so next() yields the same char.
155                        if let Some(escaped) = chars.next() {
156                            out.push(escaped);
157                        }
158                        continue;
159                    }
160                    break;
161                }
162            }
163            continue;
164        }
165        if c == '"' {
166            out.push(c);
167            for n in chars.by_ref() {
168                out.push(n);
169                if n == '"' {
170                    break;
171                }
172            }
173            continue;
174        }
175        if c.is_ascii_alphabetic() || c == '_' {
176            let mut word = String::from(c);
177            while let Some(&n) = chars.peek() {
178                if n.is_ascii_alphanumeric() || n == '_' {
179                    word.push(n);
180                    chars.next();
181                } else {
182                    break;
183                }
184            }
185            let upper = word.to_ascii_uppercase();
186            if RESERVED_FUNC_KEYWORDS.contains(&upper.as_str()) {
187                out.push_str(&word.to_ascii_lowercase());
188            } else {
189                out.push_str(&word);
190            }
191            continue;
192        }
193        out.push(c);
194    }
195    out
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use std::path::PathBuf;
202
203    fn loc() -> SourceLocation {
204        SourceLocation::new(PathBuf::from("test.sql"), 1, 1)
205    }
206
207    /// Parse the *value* expression of `SELECT <sql>` and return its `NodeEnum`.
208    fn parse_expr(sql: &str) -> NodeEnum {
209        let full = format!("SELECT {sql}");
210        let parsed = pg_query::parse(&full).expect("expression parses");
211        let stmt = parsed.protobuf.stmts.into_iter().next().expect("one stmt");
212        let select = stmt.stmt.expect("stmt").node.expect("node");
213        let NodeEnum::SelectStmt(s) = select else {
214            panic!("expected SelectStmt, got {select:?}");
215        };
216        let target = s
217            .target_list
218            .into_iter()
219            .next()
220            .expect("target")
221            .node
222            .expect("target node");
223        let NodeEnum::ResTarget(rt) = target else {
224            panic!("expected ResTarget");
225        };
226        rt.val
227            .expect("val")
228            .node
229            .expect("inner node of ResTarget val")
230    }
231
232    #[test]
233    fn cast_to_target_integer_strips() {
234        let node = parse_expr("42 :: integer");
235        let n = from_pg_node(&node, Some(&ColumnType::Integer), &loc()).expect("normalizes");
236        assert_eq!(n.canonical_text, "42");
237    }
238
239    #[test]
240    fn cast_to_target_text_strips() {
241        let node = parse_expr("'foo' :: text");
242        let n = from_pg_node(&node, Some(&ColumnType::Text), &loc()).expect("normalizes");
243        assert_eq!(n.canonical_text, "'foo'");
244    }
245
246    #[test]
247    fn cast_to_other_type_kept() {
248        let node = parse_expr("42 :: integer");
249        let n = from_pg_node(&node, Some(&ColumnType::Text), &loc()).expect("normalizes");
250        // The cast survives because the target type does not match.
251        assert!(
252            n.canonical_text.contains("integer") || n.canonical_text.contains("::"),
253            "got: {}",
254            n.canonical_text
255        );
256    }
257
258    #[test]
259    fn cast_with_no_target_kept() {
260        let node = parse_expr("42 :: integer");
261        let n = from_pg_node(&node, None, &loc()).expect("normalizes");
262        // Without a target type, no cast stripping happens.
263        assert!(
264            n.canonical_text.contains("integer") || n.canonical_text.contains("::"),
265            "got: {}",
266            n.canonical_text
267        );
268    }
269
270    #[test]
271    fn keywords_lowercased() {
272        // pg_query already lowercases most keywords; this asserts the canonical
273        // text has no uppercase reserved-word artifacts.
274        let node = parse_expr("LOWER('FOO')");
275        let n = from_pg_node(&node, None, &loc()).expect("normalizes");
276        assert!(
277            !n.canonical_text.contains("AS") && !n.canonical_text.contains("DISTINCT"),
278            "got: {}",
279            n.canonical_text
280        );
281    }
282
283    #[test]
284    fn lowercase_helper_preserves_quoted_strings() {
285        assert_eq!(lowercase_keywords("'AND OR'"), "'AND OR'");
286        assert_eq!(lowercase_keywords("foo 'AND' OR"), "foo 'AND' or");
287        assert_eq!(lowercase_keywords("AND OR"), "and or");
288    }
289
290    #[test]
291    fn lowercase_helper_handles_doubled_quote() {
292        assert_eq!(lowercase_keywords("'a''b'"), "'a''b'");
293    }
294
295    #[test]
296    fn lowercase_helper_preserves_quoted_identifiers() {
297        assert_eq!(lowercase_keywords("\"AND\""), "\"AND\"");
298    }
299
300    #[test]
301    fn hash_matches_text() {
302        let n = NormalizedExpr::from_text("now()");
303        let expected: [u8; 32] = blake3::hash(b"now()").into();
304        assert_eq!(n.ast_hash, expected);
305    }
306
307    #[test]
308    fn equivalent_casts_hash_equal() {
309        let a = parse_expr("42 :: integer");
310        let na = from_pg_node(&a, Some(&ColumnType::Integer), &loc()).unwrap();
311        let nb = NormalizedExpr::from_text("42");
312        assert_eq!(na.ast_hash, nb.ast_hash);
313    }
314}