From eb0bcb82322ec16cb5a9ca674b5bc5e5697deba4 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Mon, 28 Sep 2026 19:16:52 -0400 Subject: [PATCH] ide: add goto def & hover for ordinals --- Cargo.lock | 1 + .../src/code_actions/rewrite_integer_radix.rs | 74 +++--- crates/squawk_ide/src/collect.rs | 4 +- crates/squawk_ide/src/goto_definition.rs | 211 ++++++++++++++++++ crates/squawk_ide/src/hover.rs | 154 ++++++++++++- crates/squawk_ide/src/infer.rs | 16 +- crates/squawk_ide/src/literals.rs | 35 --- crates/squawk_ide/src/resolve.rs | 122 ++++++++++ crates/squawk_syntax/Cargo.toml | 1 + crates/squawk_syntax/src/ast.rs | 2 +- crates/squawk_syntax/src/ast/node_ext.rs | 82 +++++++ 11 files changed, 612 insertions(+), 90 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c176f8571..2079546a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2578,6 +2578,7 @@ dependencies = [ "dir-test", "either", "insta", + "num-bigint", "rowan", "smol_str", "squawk-line-index", diff --git a/crates/squawk_ide/src/code_actions/rewrite_integer_radix.rs b/crates/squawk_ide/src/code_actions/rewrite_integer_radix.rs index f85dcab3d..ccdb16686 100644 --- a/crates/squawk_ide/src/code_actions/rewrite_integer_radix.rs +++ b/crates/squawk_ide/src/code_actions/rewrite_integer_radix.rs @@ -2,40 +2,42 @@ use num_bigint::BigUint; use rowan::TextSize; use salsa::Database as Db; use squawk_linter::Edit; -use squawk_syntax::SyntaxKind; - -use crate::{ - file::InFile, - literals::{IntegerRadix, normalize_integer_literal}, - offsets::token_from_offset, +use squawk_syntax::{ + SyntaxKind, + ast::{self, AstNode}, }; -use super::{ActionKind, CodeAction}; +use crate::{file::InFile, offsets::token_from_offset}; -impl IntegerRadix { - const ALL: [Self; 4] = [Self::Binary, Self::Octal, Self::Decimal, Self::Hexadecimal]; +use super::{ActionKind, CodeAction}; - fn format(self, value: &BigUint) -> String { - let mut digits = value.to_str_radix(self.base()); - if self == Self::Hexadecimal { - digits.make_ascii_uppercase(); - } - let prefix = match self { - Self::Binary => "0b", - Self::Decimal => "", - Self::Hexadecimal => "0x", - Self::Octal => "0o", - }; - format!("{prefix}{digits}") +const INTEGER_RADIXES: [ast::IntegerRadix; 4] = [ + ast::IntegerRadix::Binary, + ast::IntegerRadix::Octal, + ast::IntegerRadix::Decimal, + ast::IntegerRadix::Hexadecimal, +]; + +fn format_integer(radix: ast::IntegerRadix, value: &BigUint) -> String { + let mut digits = value.to_str_radix(radix.base()); + if radix == ast::IntegerRadix::Hexadecimal { + digits.make_ascii_uppercase(); } + let prefix = match radix { + ast::IntegerRadix::Binary => "0b", + ast::IntegerRadix::Decimal => "", + ast::IntegerRadix::Hexadecimal => "0x", + ast::IntegerRadix::Octal => "0o", + }; + format!("{prefix}{digits}") +} - fn name(self) -> &'static str { - match self { - Self::Binary => "binary", - Self::Decimal => "decimal", - Self::Hexadecimal => "hexadecimal", - Self::Octal => "octal", - } +fn integer_radix_name(radix: ast::IntegerRadix) -> &'static str { + match radix { + ast::IntegerRadix::Binary => "binary", + ast::IntegerRadix::Decimal => "decimal", + ast::IntegerRadix::Hexadecimal => "hexadecimal", + ast::IntegerRadix::Octal => "octal", } } @@ -49,15 +51,17 @@ pub(super) fn rewrite_integer_radix( return None; } - let (source_radix, value) = parse_integer_literal(token.text())?; - for target_radix in IntegerRadix::ALL { + let literal = token.parent().and_then(ast::Literal::cast)?; + let source_radix = literal.integer_radix()?; + let value = literal.integer_value()?; + for target_radix in INTEGER_RADIXES { if target_radix == source_radix { continue; } - let replacement = target_radix.format(&value); + let replacement = format_integer(target_radix, &value); actions.push(CodeAction { - title: format!("Rewrite integer as {}", target_radix.name()), + title: format!("Rewrite integer as {}", integer_radix_name(target_radix)), edits: vec![Edit::replace(token.text_range(), replacement)], kind: ActionKind::RefactorRewrite, }); @@ -66,12 +70,6 @@ pub(super) fn rewrite_integer_radix( Some(()) } -fn parse_integer_literal(text: &str) -> Option<(IntegerRadix, BigUint)> { - let (radix, digits) = normalize_integer_literal(text); - let value = BigUint::parse_bytes(digits.as_bytes(), radix.base())?; - Some((radix, value)) -} - #[cfg(test)] mod tests { use insta::assert_snapshot; diff --git a/crates/squawk_ide/src/collect.rs b/crates/squawk_ide/src/collect.rs index 6a63d6adb..a85823175 100644 --- a/crates/squawk_ide/src/collect.rs +++ b/crates/squawk_ide/src/collect.rs @@ -542,7 +542,7 @@ fn with_table_query_columns_with_types( target_list_columns_with_types_in_file(db, file, &target_list, from_clause.as_ref()) } -fn target_expr_type(db: &dyn Db, file: FileId, target: &ast::Target) -> Option { +pub(crate) fn target_expr_type(db: &dyn Db, file: FileId, target: &ast::Target) -> Option { let expr = target.expr()?; if let Some(ty) = infer_type_from_expr(&expr) { return Some(ty); @@ -653,7 +653,7 @@ fn columns_for_star_from_clause( columns } -fn columns_for_star_from_from_item( +pub(crate) fn columns_for_star_from_from_item( db: &dyn Db, file: FileId, from_item: &ast::FromItem, diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 5bbf3c641..bcd232ddc 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -8592,6 +8592,217 @@ select * from t group by t.b$0; "); } + #[test] + fn goto_select_target_ordinal_in_group_by() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b) +select a, max(b) from t group by 1$0 +"), @" + ╭▸ + 3 │ select a, max(b) from t group by 1 + ╰╴ ─ 2. destination ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_order_by() { + assert_snapshot!(goto(" +select a, max(b) as maximum from t order by 2$0 +"), @" + ╭▸ + 2 │ select a, max(b) as maximum from t order by 2 + ╰╴ ───────────────── 2. destination ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_parenthesized() { + assert_snapshot!(goto(" +select a from t order by ((1$0)) +"), @" + ╭▸ + 2 │ select a from t order by ((1)) + ╰╴ ─ 2. destination ─ 1. source + "); + assert_snapshot!(goto(" +select a from t group by ((1$0)) +"), @" + ╭▸ + 2 │ select a from t group by ((1)) + ╰╴ ─ 2. destination ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_distinct_on() { + assert_snapshot!(goto(" +select distinct on (1$0) a from t +"), @" + ╭▸ + 2 │ select distinct on (1) a from t + │ ┬ ─ 2. destination + │ │ + ╰╴ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_group_by_set() { + assert_snapshot!(goto(" +select a from t group by cube (1$0) +"), @" + ╭▸ + 2 │ select a from t group by cube (1) + ╰╴ ─ 2. destination ─ 1. source + "); + assert_snapshot!(goto(" +select a from t group by rollup (1$0) +"), @" + ╭▸ + 2 │ select a from t group by rollup (1) + ╰╴ ─ 2. destination ─ 1. source + "); + assert_snapshot!(goto(" +select a from t group by grouping sets (1$0) +"), @" + ╭▸ + 2 │ select a from t group by grouping sets (1) + ╰╴ ─ 2. destination ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_out_of_range() { + goto_not_found("select a from t order by 2$0"); + } + + #[test] + fn goto_select_target_ordinal_in_star() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b) +select *, a from t order by 2$0 +"), @" + ╭▸ + 2 │ with t as (select 1 a, 2 b) + │ ─ 2. destination + 3 │ select *, a from t order by 2 + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_after_star() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b) +select *, a from t order by 3$0 +"), @" + ╭▸ + 3 │ select *, a from t order by 3 + ╰╴ ─ 2. destination ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_qualified_star() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b), + u as (select 3 c, 4 d) +select u.*, t.* from t, u group by 3$0 +"), @" + ╭▸ + 2 │ with t as (select 1 a, 2 b), + │ ─ 2. destination + 3 │ u as (select 3 c, 4 d) + 4 │ select u.*, t.* from t, u group by 3 + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_star_across_from_items() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b), + u as (select 3 c, 4 d) +select * from t, u order by 4$0 +"), @" + ╭▸ + 3 │ u as (select 3 c, 4 d) + │ ─ 2. destination + 4 │ select * from t, u order by 4 + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_star_through_cte_star() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b), + u as (select * from t) +select * from u order by 2$0 +"), @" + ╭▸ + 2 │ with t as (select 1 a, 2 b), + │ ─ 2. destination + 3 │ u as (select * from t) + 4 │ select * from u order by 2 + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_star_with_duplicate_names() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 a) +select * from t order by 2$0 +"), @" + ╭▸ + 2 │ with t as (select 1 a, 2 a) + │ ─ 2. destination + 3 │ select * from t order by 2 + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_select_target_ordinal_in_star_with_alias_columns() { + assert_snapshot!(goto(" +with t as (select 1 a, 2 b) +select * from t as x(y) order by 1$0 +"), @" + ╭▸ + 3 │ select * from t as x(y) order by 1 + │ ┬ ─ 1. source + │ │ + ╰╴ 2. destination + "); + } + + #[test] + fn goto_select_target_ordinal_in_star_out_of_range() { + goto_not_found( + " +with t as (select 1 a, 2 b) +select *, a from t order by 4$0 +", + ); + } + + #[test] + fn goto_select_target_ordinal_in_paren_select() { + assert_snapshot!(goto(" +(select a, b from t) order by 2$0 +"), @" + ╭▸ + 2 │ (select a, b from t) order by 2 + ╰╴ ─ 2. destination ─ 1. source + "); + } + + #[test] + fn goto_window_order_by_integer_is_not_a_target_ordinal() { + goto_not_found("select a, row_number() over (order by 1$0) from t"); + } + #[test] fn goto_select_alias_order_by_column_name_conflict() { // If an ORDER BY expression is a simple name that matches both an diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 8b3c58f83..f55230a1d 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -259,12 +259,37 @@ pub fn hover(db: &dyn Db, position: InFile) -> Option { } if let Some(literal) = ast::Literal::cast(parent) { + if literal.integer_value().is_some() { + return hover_select_target_ordinal(db, position); + } return hover_literal(&literal); } None } +fn hover_select_target_ordinal(db: &dyn Db, position: InFile) -> Option { + let def = *goto_definition::goto_definition(db, position).first()?; + if let Some(target) = def.to_node(db)?.ancestors().find_map(ast::Target::cast) + && let Some((_, node)) = ColumnName::from_target(target) + { + return hover(db, InFile::new(def.file, node.text_range().start())); + } + hover_position(db, position) +} + +fn hover_select_target(db: &dyn Db, def: Location) -> Option { + let target = def.to_node(db)?.ancestors().find_map(ast::Target::cast)?; + let (column_name, _) = ColumnName::from_target(target.clone())?; + let column_name = column_name.to_string()?; + Some(Hover::snippet( + match collect::target_expr_type(db, def.file, &target) { + Some(ty) => ColumnHover::anon_column_type(&column_name, &ty.to_string()), + None => ColumnHover::anon_column(&column_name), + }, + )) +} + fn hover_literal(literal: &ast::Literal) -> Option { let kind = literal.kind()?; // TODO: support all literal types @@ -469,7 +494,10 @@ fn hover_position(db: &dyn Db, position: InFile) -> Option { return Some(result); } // Finally try as table (handles case like `select t from t;` where t is the table) - hover_table(db, def) + if let Some(result) = hover_table(db, def) { + return Some(result); + } + hover_select_target(db, def) } LocationKind::Collation => hover_collation(db, def), LocationKind::Constraint => hover_constraint(db, def), @@ -2191,6 +2219,13 @@ mod test { None } + #[track_caller] + fn hover_not_found(sql: &str) { + if let Some(hover) = check_hover_(sql) { + panic!("expected no hover, found:\n{hover}"); + } + } + #[must_use] #[track_caller] fn check_hover_info(sql: &str) -> super::Hover { @@ -2200,6 +2235,116 @@ mod test { hover(fixture.db(), offset).expect("should find hover information") } + #[test] + fn hover_select_target_ordinal_in_order_by() { + assert_snapshot!(check_hover(" +create table t(a int, b text); +select a, b from t order by 2$0; +"), @" + hover: column public.t.b text + ╭▸ + 3 │ select a, b from t order by 2; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_target_ordinal_in_group_by() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b) +select a, max(b) from t group by 1$0; +"), @" + hover: column t.a integer + ╭▸ + 3 │ select a, max(b) from t group by 1; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_target_ordinal_aliased_expr() { + assert_snapshot!(check_hover(" +create table t(a int, b int); +select a, max(b) as maximum from t order by 2$0; +"), @" + hover: column maximum + ╭▸ + 3 │ select a, max(b) as maximum from t order by 2; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_target_ordinal_aliased_cast() { + assert_snapshot!(check_hover(" +create table t(a int); +select a::text as label from t order by 1$0; +"), @" + hover: column label text + ╭▸ + 3 │ select a::text as label from t order by 1; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_window_order_by_integer_is_not_a_target_ordinal() { + hover_not_found("select a, row_number() over (order by 1$0) from t"); + } + + #[test] + fn hover_select_target_alias() { + assert_snapshot!(check_hover(" +create table t(a int); +select a::text as label$0 from t; +"), @" + hover: column label text + ╭▸ + 3 │ select a::text as label from t; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_target_ordinal_in_star() { + assert_snapshot!(check_hover(" +with t as (select 1 a, 2 b) +select *, a from t order by 2$0; +"), @" + hover: column t.b integer + ╭▸ + 3 │ select *, a from t order by 2; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_target_ordinal_in_qualified_star() { + assert_snapshot!(check_hover(" +create table t(a int, b text); +create table u(c bigint); +select u.*, t.* from t, u order by 3$0; +"), @" + hover: column public.t.b text + ╭▸ + 4 │ select u.*, t.* from t, u order by 3; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_target_ordinal_in_paren_select() { + assert_snapshot!(check_hover(" +create table t(a int, b text); +(select a, b from t) order by 2$0; +"), @" + hover: column public.t.b text + ╭▸ + 3 │ (select a, b from t) order by 2; + ╰╴ ─ hover + "); + } + #[test] fn hover_column_in_create_index() { assert_snapshot!(check_hover(" @@ -4076,14 +4221,11 @@ select u.*$0 from t u(x, y); #[test] fn hover_on_star_from_cte_empty_select() { - assert!( - check_hover_( - " + hover_not_found( + " with t as (select) select *$0 from t; ", - ) - .is_none() ); } diff --git a/crates/squawk_ide/src/infer.rs b/crates/squawk_ide/src/infer.rs index ad8789efa..d8381bfab 100644 --- a/crates/squawk_ide/src/infer.rs +++ b/crates/squawk_ide/src/infer.rs @@ -5,8 +5,6 @@ use squawk_syntax::{ ast::{self, AstNode}, }; -use crate::literals::normalize_integer_literal; - #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum Type { Array(Box), @@ -69,11 +67,13 @@ pub(crate) fn infer_type_from_ty(ty: &ast::Type) -> Option { } } -fn infer_int_type(text: &str) -> Type { - let (radix, digits) = normalize_integer_literal(text); - match u64::from_str_radix(&digits, radix.base()) { - Ok(n) if n <= i32::MAX as u64 => Type::Integer, - Ok(n) if n <= i64::MAX as u64 => Type::Bigint, +fn infer_int_type(literal: &ast::Literal) -> Type { + match literal + .integer_value() + .and_then(|value| u64::try_from(value).ok()) + { + Some(n) if n <= i32::MAX as u64 => Type::Integer, + Some(n) if n <= i64::MAX as u64 => Type::Bigint, _ => Type::Numeric, } } @@ -81,7 +81,7 @@ fn infer_int_type(text: &str) -> Type { pub(crate) fn infer_type_from_literal(literal: &ast::Literal) -> Option { let token = literal.syntax().first_token()?; match token.kind() { - SyntaxKind::INT_NUMBER => Some(infer_int_type(token.text())), + SyntaxKind::INT_NUMBER => Some(infer_int_type(literal)), SyntaxKind::NUMERIC_NUMBER => Some(Type::Numeric), // TODO: this isn't necessarily text, e.g., select 1 + '1'; // We need to look at the context of the string's usage to be sure. diff --git a/crates/squawk_ide/src/literals.rs b/crates/squawk_ide/src/literals.rs index b433defb2..8edcb51a5 100644 --- a/crates/squawk_ide/src/literals.rs +++ b/crates/squawk_ide/src/literals.rs @@ -1,38 +1,3 @@ -#[derive(Clone, Copy, PartialEq, Eq)] -pub(crate) enum IntegerRadix { - Binary, - Decimal, - Hexadecimal, - Octal, -} - -impl IntegerRadix { - pub(crate) fn base(self) -> u32 { - match self { - Self::Binary => 2, - Self::Decimal => 10, - Self::Hexadecimal => 16, - Self::Octal => 8, - } - } -} - -pub(crate) fn normalize_integer_literal(text: &str) -> (IntegerRadix, String) { - let (radix, digits) = match text.as_bytes() { - [b'0', b'b' | b'B', ..] => (IntegerRadix::Binary, &text[2..]), - [b'0', b'o' | b'O', ..] => (IntegerRadix::Octal, &text[2..]), - [b'0', b'x' | b'X', ..] => (IntegerRadix::Hexadecimal, &text[2..]), - _ => (IntegerRadix::Decimal, text), - }; - let digits = if radix == IntegerRadix::Decimal { - digits - } else { - digits.strip_prefix('_').unwrap_or(digits) - }; - - (radix, digits.replace('_', "")) -} - pub(crate) fn binary_digits_to_hex(digits: &str) -> Option { const HEX_DIGITS: &[u8; 16] = b"0123456789ABCDEF"; diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index d1597fcbb..f49e7324a 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -10,6 +10,7 @@ use squawk_syntax::{ }; use crate::binder::ResolvedSchemas; +use crate::collect::columns_for_star_from_from_item; use crate::db::FileId; use crate::file::InFile; use crate::location::{Location, LocationKind}; @@ -988,6 +989,10 @@ pub(crate) fn resolve_literal( return resolve_positional_param(db, InFile::new(file, literal), index); } + if let Some(locations) = resolve_select_target_ordinal(db, file, literal) { + return Some(locations); + } + let context = classify_literal(literal.syntax())?; match context { @@ -1017,6 +1022,123 @@ pub(crate) fn resolve_literal( } } +fn resolve_select_target_ordinal( + db: &dyn Db, + file: FileId, + literal: &ast::Literal, +) -> Option> { + let ordinal = usize::try_from(literal.integer_value()?) + .ok()? + .checked_sub(1)?; + let owner = select_target_ordinal_owner(literal)?; + let (target_list, from_clause) = targets_for_select_variant(ast::SelectVariant::cast(owner)?)?; + + let mut remaining = ordinal; + for target in target_list.targets() { + if !target_is_star(&target) { + if remaining == 0 { + return Some(smallvec![Location::new( + file, + target.syntax().text_range(), + LocationKind::Column, + )]); + } + remaining -= 1; + continue; + } + + for from_item in star_target_from_items(&target, from_clause.as_ref()?)? { + let columns = columns_for_star_from_from_item(db, file, &from_item); + if columns.is_empty() { + return None; + } + if let Some((column_name, _)) = columns.get(remaining) { + let scope_name_ref = relation_name_ref_from_from_item(&from_item)?; + return resolve_from_item_column_by_name_after_index( + db, + InFile::new(file, &from_item), + &scope_name_ref, + column_name, + remaining, + ); + } + remaining -= columns.len(); + } + } + None +} + +fn star_target_from_items( + target: &ast::Target, + from_clause: &ast::FromClause, +) -> Option> { + if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr() { + let table_name = qualified_star_table_name(&field_expr)?; + return Some(vec![find_from_item_in_from_clause( + from_clause, + &table_name, + )?]); + } + Some(ast_nav::iter_from_clause(from_clause).collect()) +} + +fn select_target_ordinal_owner(literal: &ast::Literal) -> Option { + let parent = literal + .syntax() + .ancestors() + .skip(1) + .find(|node| !ast::ParenExpr::can_cast(node.kind()))?; + if ast::SortBy::can_cast(parent.kind()) { + return parent + .ancestors() + .find_map(ast::OrderByClause::cast)? + .syntax() + .parent(); + } + if ast::GroupingExpr::can_cast(parent.kind()) + || ast::GroupingCube::can_cast(parent.kind()) + || ast::GroupingRollup::can_cast(parent.kind()) + { + return parent + .ancestors() + .find_map(ast::GroupByClause::cast)? + .syntax() + .parent(); + } + if ast::DistinctOn::can_cast(parent.kind()) { + return parent + .ancestors() + .find_map(ast::SelectClause::cast)? + .syntax() + .parent(); + } + None +} + +fn target_is_star(target: &ast::Target) -> bool { + target.star_token().is_some() + || matches!( + target.expr(), + Some(ast::Expr::FieldExpr(field_expr)) if field_expr.star_token().is_some() + ) +} + +fn targets_for_select_variant( + select: ast::SelectVariant, +) -> Option<(ast::TargetList, Option)> { + match select { + ast::SelectVariant::CompoundSelect(select) => targets_for_select_variant(select.lhs()?), + ast::SelectVariant::ParenSelect(select) => targets_for_select_variant(select.select()?), + ast::SelectVariant::Select(select) => { + Some((select.select_clause()?.target_list()?, select.from_clause())) + } + ast::SelectVariant::SelectInto(select) => { + Some((select.select_clause()?.target_list()?, select.from_clause())) + } + ast::SelectVariant::Table(_) | ast::SelectVariant::Values(_) => None, + } +} + pub(crate) fn resolve_custom_op( db: &dyn Db, custom_op: InFile<&ast::CustomOp>, diff --git a/crates/squawk_syntax/Cargo.toml b/crates/squawk_syntax/Cargo.toml index 7574e79b0..657977fcc 100644 --- a/crates/squawk_syntax/Cargo.toml +++ b/crates/squawk_syntax/Cargo.toml @@ -19,6 +19,7 @@ squawk-parser.workspace = true rowan.workspace = true smol_str.workspace = true either.workspace = true +num-bigint.workspace = true [dev-dependencies] annotate-snippets.workspace = true diff --git a/crates/squawk_syntax/src/ast.rs b/crates/squawk_syntax/src/ast.rs index 9c4acc0f1..7216477d1 100644 --- a/crates/squawk_syntax/src/ast.rs +++ b/crates/squawk_syntax/src/ast.rs @@ -39,7 +39,7 @@ use squawk_parser::SyntaxKind; pub use self::node_ext::{LanguageName, is_quoted_name_node, normalize_name_node}; pub use self::{ generated::tokens::*, - node_ext::{BinOp, CastKind, LitKind, PostfixOp, PrefixOp}, + node_ext::{BinOp, CastKind, IntegerRadix, LitKind, PostfixOp, PrefixOp}, nodes::*, traits::{HasCreateTable, HasPathRef, HasSelectTail, HasWithClause, NameLike}, }; diff --git a/crates/squawk_syntax/src/ast/node_ext.rs b/crates/squawk_syntax/src/ast/node_ext.rs index 082750327..f4978ad46 100644 --- a/crates/squawk_syntax/src/ast/node_ext.rs +++ b/crates/squawk_syntax/src/ast/node_ext.rs @@ -29,6 +29,7 @@ use std::borrow::Cow; use either::Either; #[cfg(test)] use insta::assert_snapshot; +use num_bigint::BigUint; use rowan::{GreenNodeData, GreenTokenData, NodeOrToken, TextRange, TextSize}; use squawk_line_index::{LineEnding, find_newline}; @@ -89,6 +90,25 @@ pub enum CastKind { TypeLiteral, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IntegerRadix { + Binary, + Decimal, + Hexadecimal, + Octal, +} + +impl IntegerRadix { + pub fn base(self) -> u32 { + match self { + Self::Binary => 2, + Self::Decimal => 10, + Self::Hexadecimal => 16, + Self::Octal => 8, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum LitKind { BitString(SyntaxToken), @@ -211,6 +231,31 @@ impl ast::Literal { .ok()? .checked_sub(1) } + + pub fn integer_radix(&self) -> Option { + let LitKind::IntNumber(token) = self.kind()? else { + return None; + }; + Some(match token.text().as_bytes() { + [b'0', b'b' | b'B', ..] => IntegerRadix::Binary, + [b'0', b'o' | b'O', ..] => IntegerRadix::Octal, + [b'0', b'x' | b'X', ..] => IntegerRadix::Hexadecimal, + _ => IntegerRadix::Decimal, + }) + } + + pub fn integer_value(&self) -> Option { + let LitKind::IntNumber(token) = self.kind()? else { + return None; + }; + let radix = self.integer_radix()?; + let text = token.text(); + let digits = match radix { + IntegerRadix::Decimal => text, + IntegerRadix::Binary | IntegerRadix::Hexadecimal | IntegerRadix::Octal => &text[2..], + }; + BigUint::parse_bytes(digits.replace('_', "").as_bytes(), radix.base()) + } } #[derive(Clone, Copy, PartialEq, Eq)] @@ -2185,6 +2230,43 @@ fn vacuum_full_0_is_not_full() { assert!(!extract_vacuum("VACUUM (FULL 0) foo;").is_full()); } +#[cfg(test)] +fn extract_literal(sql: &str) -> ast::Literal { + let parse = SourceFile::parse(sql); + assert!(parse.errors().is_empty(), "{:?}", parse.errors()); + parse + .tree() + .syntax() + .descendants() + .find_map(ast::Literal::cast) + .unwrap() +} + +#[test] +fn integer_value() { + assert_eq!( + extract_literal("select 42").integer_value(), + Some(42u8.into()) + ); + assert_eq!( + extract_literal("select 1_000").integer_value(), + Some(1_000u16.into()) + ); + assert_eq!( + extract_literal("select 0b1010").integer_value(), + Some(10u8.into()) + ); + assert_eq!( + extract_literal("select 0o12").integer_value(), + Some(10u8.into()) + ); + assert_eq!( + extract_literal("select 0xA").integer_value(), + Some(10u8.into()) + ); + assert_eq!(extract_literal("select 1.0").integer_value(), None); +} + #[cfg(test)] fn decode_literal(sql: &str) -> String { let parse = SourceFile::parse(sql);