ide: refactor resolve to use name related functions

Cargo.toml

@@ -1,5 +1,6 @@
 [workspace]
 members = ["crates/*"]
+default-members = ["crates/squawk"]
 resolver = "2"
 
 [workspace.package]

crates/squawk_ide/src/ast_nav.rs

@@ -2,9 +2,10 @@
 ///
 /// There shouldn't be any dependency on Salsa.
 use squawk_syntax::{
-    SyntaxNode,
+    SyntaxNode, SyntaxToken,
     ast::{self, AstNode},
 };
+use std::iter;
 
 use crate::symbols::Name;
 
@@ -74,6 +75,71 @@ pub(crate) fn node_parent_query(node: &SyntaxNode) -> Option<ParentQuery> {
     None
 }
 
+#[derive(Debug)]
+pub(crate) enum SelectContext {
+    Compound(ast::CompoundSelect),
+    Single(ast::Select),
+}
+
+impl SelectContext {
+    pub(crate) fn iter(&self) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
+        fn variant_iter(
+            variant: ast::SelectVariant,
+        ) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
+            match variant {
+                ast::SelectVariant::Select(select) => Some(Box::new(iter::once(select))),
+                ast::SelectVariant::CompoundSelect(compound) => compound_iter(&compound),
+                ast::SelectVariant::ParenSelect(_)
+                | ast::SelectVariant::SelectInto(_)
+                | ast::SelectVariant::Table(_)
+                | ast::SelectVariant::Values(_) => None,
+            }
+        }
+
+        fn compound_iter(
+            node: &ast::CompoundSelect,
+        ) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
+            let lhs_iter = node
+                .lhs()
+                .map(variant_iter)
+                .unwrap_or_else(|| Some(Box::new(iter::empty())))?;
+            let rhs_iter = node
+                .rhs()
+                .map(variant_iter)
+                .unwrap_or_else(|| Some(Box::new(iter::empty())))?;
+            Some(Box::new(lhs_iter.chain(rhs_iter)))
+        }
+
+        match self {
+            SelectContext::Compound(compound) => compound_iter(compound),
+            SelectContext::Single(select) => Some(Box::new(iter::once(select.clone()))),
+        }
+    }
+}
+
+pub(crate) fn find_select_parent(token: SyntaxToken) -> Option<SelectContext> {
+    let mut found_select = None;
+    let mut found_compound = None;
+
+    for ancestor in token.parent_ancestors() {
+        if let Some(compound_select) = ast::CompoundSelect::cast(ancestor.clone()) {
+            if compound_select.union_token().is_some() && compound_select.all_token().is_some() {
+                found_compound = Some(SelectContext::Compound(compound_select));
+            } else {
+                break;
+            }
+        }
+
+        if found_select.is_none()
+            && let Some(select) = ast::Select::cast(ancestor)
+        {
+            found_select = Some(SelectContext::Single(select));
+        }
+    }
+
+    found_compound.or(found_select)
+}
+
 ///
 /// ```sql
 /// with t as (select 1)
@@ -164,6 +230,28 @@ pub(crate) fn parent_source(node: &SyntaxNode) -> Option<ParentSouce> {
     None
 }
 
+struct UnwrapParenExpr {
+    current: Option<ast::Expr>,
+}
+
+impl Iterator for UnwrapParenExpr {
+    type Item = ast::Expr;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let expr = self.current.take()?;
+        if let ast::Expr::ParenExpr(paren_expr) = &expr {
+            self.current = paren_expr.expr();
+        }
+        Some(expr)
+    }
+}
+
+pub(crate) fn unwrap_paren_expr(expr: ast::Expr) -> impl Iterator<Item = ast::Expr> {
+    UnwrapParenExpr {
+        current: Some(expr),
+    }
+}
+
 pub(crate) fn iter_from_clause(
     from_clause: &ast::FromClause,
 ) -> impl Iterator<Item = ast::FromItem> {

crates/squawk_ide/src/binder.rs

@@ -132,16 +132,15 @@ impl Binder {
 
     pub(crate) fn lookup_info(
         &self,
-        name_str: String,
-        schema: &Option<String>,
+        name: &Name,
+        schema: Option<&Schema>,
         kind: SymbolKind,
         position: TextSize,
     ) -> Option<(Schema, String)> {
-        let name_normalized = Name::from_string(name_str.clone());
-        let symbols = self.scope.get(&name_normalized)?;
+        let symbols = self.scope.get(name)?;
 
-        let search_paths = match schema {
-            Some(schema_name) => &[Schema::new(schema_name)],
+        let search_paths: &[Schema] = match schema {
+            Some(s) => std::slice::from_ref(s),
             None => self.search_path_at(position),
         };
 
@@ -151,7 +150,7 @@ impl Binder {
                 symbol.kind == kind && symbol.schema.as_ref() == Some(search_schema)
             }) {
                 let symbol = &self.symbols[symbol_id];
-                return Some((symbol.schema.clone()?, name_str));
+                return Some((symbol.schema.clone()?, name.to_string()));
             }
         }
         None

crates/squawk_ide/src/code_actions/rewrite_select_as_values.rs

@@ -1,12 +1,13 @@
 use rowan::{TextRange, TextSize};
 use salsa::Database as Db;
-use squawk_syntax::{
-    SyntaxToken,
-    ast::{self, AstNode},
-};
-use std::iter;
+use squawk_syntax::ast::{self, AstNode};
 
-use crate::{db::File, offsets::token_from_offset, symbols::Name};
+use crate::{
+    ast_nav::{self, SelectContext},
+    db::File,
+    offsets::token_from_offset,
+    symbols::Name,
+};
 
 use super::{ActionKind, CodeAction};
 
@@ -18,7 +19,7 @@ pub(super) fn rewrite_select_as_values(
 ) -> Option<()> {
     let token = token_from_offset(db, file, offset)?;
 
-    let parent = find_select_parent(token)?;
+    let parent = ast_nav::find_select_parent(token)?;
 
     let mut selects = parent.iter()?.peekable();
     let select_token_start = selects
@@ -81,67 +82,6 @@ fn is_values_row_column_name(target: &ast::Target, idx: usize) -> bool {
     true
 }
 
-enum SelectContext {
-    Compound(ast::CompoundSelect),
-    Single(ast::Select),
-}
-
-impl SelectContext {
-    fn iter(&self) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
-        fn variant_iter(
-            variant: ast::SelectVariant,
-        ) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
-            match variant {
-                ast::SelectVariant::Select(select) => Some(Box::new(iter::once(select))),
-                ast::SelectVariant::CompoundSelect(compound) => compound_iter(&compound),
-                ast::SelectVariant::ParenSelect(_)
-                | ast::SelectVariant::SelectInto(_)
-                | ast::SelectVariant::Table(_)
-                | ast::SelectVariant::Values(_) => None,
-            }
-        }
-
-        fn compound_iter(
-            node: &ast::CompoundSelect,
-        ) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
-            let lhs_iter = node
-                .lhs()
-                .map(variant_iter)
-                .unwrap_or_else(|| Some(Box::new(iter::empty())))?;
-            let rhs_iter = node
-                .rhs()
-                .map(variant_iter)
-                .unwrap_or_else(|| Some(Box::new(iter::empty())))?;
-            Some(Box::new(lhs_iter.chain(rhs_iter)))
-        }
-
-        match self {
-            SelectContext::Compound(compound) => compound_iter(compound),
-            SelectContext::Single(select) => Some(Box::new(iter::once(select.clone()))),
-        }
-    }
-}
-
-fn find_select_parent(token: SyntaxToken) -> Option<SelectContext> {
-    let mut found_select = None;
-    let mut found_compound = None;
-    for node in token.parent_ancestors() {
-        if let Some(compound_select) = ast::CompoundSelect::cast(node.clone()) {
-            if compound_select.union_token().is_some() && compound_select.all_token().is_some() {
-                found_compound = Some(SelectContext::Compound(compound_select));
-            } else {
-                break;
-            }
-        }
-        if found_select.is_none()
-            && let Some(select) = ast::Select::cast(node)
-        {
-            found_select = Some(SelectContext::Single(select));
-        }
-    }
-    found_compound.or(found_select)
-}
-
 #[cfg(test)]
 mod test {
     use insta::assert_snapshot;

crates/squawk_ide/src/goto_definition.rs

@@ -4629,6 +4629,23 @@ select u.id$0 from u;
         ");
     }
 
+    #[test]
+    fn goto_qualified_table_ref_prefers_schema_qualified_from_item_over_cte() {
+        assert_snapshot!(goto("
+create schema s;
+create table s.t(a int);
+with t as (select 1 a)
+select t$0.a from s.t;
+"), @r"
+          ╭▸ 
+        3 │ create table s.t(a int);
+          │                ─ 2. destination
+        4 │ with t as (select 1 a)
+        5 │ select t.a from s.t;
+          ╰╴       ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_subquery_qualified_column() {
         assert_snapshot!(goto("

crates/squawk_ide/src/location.rs

@@ -39,12 +39,6 @@ pub enum LocationKind {
     Window,
 }
 
-impl LocationKind {
-    pub(crate) fn from_node(node: &SyntaxNode) -> Option<LocationKind> {
-        classify_def_node(node)
-    }
-}
-
 #[derive(Clone, Copy, PartialEq, Eq)]
 pub struct Location {
     pub file: File,
@@ -58,7 +52,7 @@ impl Location {
     }
 
     pub(crate) fn from_node(file: File, node: &SyntaxNode) -> Option<Location> {
-        let kind = LocationKind::from_node(node)?;
+        let kind = classify_def_node(node)?;
         Some(Location::new(file, node.text_range(), kind))
     }
 

crates/squawk_ide/src/name.rs

@@ -56,8 +56,7 @@ pub(crate) fn schema_and_table_name(name_ref: &ast::NameRef) -> Option<(Option<S
 pub(crate) fn schema_and_name(name_ref: &ast::NameRef) -> (Option<Schema>, Name) {
     let table_name = Name::from_node(name_ref);
     let schema = if let Some(parent) = name_ref.syntax().parent()
-        && let Some(field_expr) = ast::FieldExpr::cast(parent)
-        && let Some(base) = field_expr.base()
+        && let Some(base) = ast::FieldExpr::cast(parent).and_then(|x| x.base())
         && let Some(schema_name_ref) = ast::NameRef::cast(base.syntax().clone())
     {
         Some(Schema(Name::from_node(&schema_name_ref)))
@@ -85,8 +84,13 @@ pub(crate) fn schema_and_func_name(call_expr: &ast::CallExpr) -> Option<(Option<
 
 pub(crate) fn table_name(path: &ast::Path) -> Option<Name> {
     let segment = path.segment()?;
-    let name_ref = segment.name_ref()?;
-    Some(Name::from_node(&name_ref))
+    if let Some(name_ref) = segment.name_ref() {
+        return Some(Name::from_node(&name_ref));
+    }
+    if let Some(name) = segment.name() {
+        return Some(Name::from_node(&name));
+    }
+    None
 }
 
 pub(crate) fn schema_name(path: &ast::Path) -> Option<Schema> {
@@ -96,6 +100,42 @@ pub(crate) fn schema_name(path: &ast::Path) -> Option<Schema> {
         .map(|name_ref| Schema(Name::from_node(&name_ref)))
 }
 
+pub(crate) fn schema_and_table_from_from_item(
+    from_item: &ast::FromItem,
+) -> Option<(Option<Schema>, Name)> {
+    if let Some(name_ref_node) = from_item.name_ref() {
+        Some((None, Name::from_node(&name_ref_node)))
+    } else if let Some(from_field_expr) = from_item.field_expr() {
+        let table_name = Name::from_node(&from_field_expr.field()?);
+        let ast::Expr::NameRef(schema_name_ref) = from_field_expr.base()? else {
+            return None;
+        };
+        let schema = Schema(Name::from_node(&schema_name_ref));
+        Some((Some(schema), table_name))
+    } else {
+        None
+    }
+}
+
+pub(crate) fn schema_and_table_from_field_expr(
+    field_expr: &ast::FieldExpr,
+) -> Option<(Option<Schema>, Name)> {
+    match field_expr.base()? {
+        ast::Expr::NameRef(name_ref) => Some((None, Name::from_node(&name_ref))),
+        ast::Expr::FieldExpr(field_expr) => {
+            let field = field_expr.field()?;
+            let ast::Expr::NameRef(schema) = field_expr.base()? else {
+                return None;
+            };
+            Some((
+                Some(Schema(Name::from_node(&schema))),
+                Name::from_node(&field),
+            ))
+        }
+        _ => None,
+    }
+}
+
 pub(crate) fn schema_and_type_name(ty: &ast::Type) -> Option<(Option<Schema>, Name)> {
     match ty {
         ast::Type::ArrayType(array_type) => {
@@ -116,8 +156,7 @@ pub(crate) fn schema_and_type_name(ty: &ast::Type) -> Option<(Option<Schema>, Na
             schema_and_name_path(&path)
         }
         ast::Type::ExprType(expr_type) => {
-            let expr = expr_type.expr()?;
-            if let ast::Expr::FieldExpr(field_expr) = expr
+            if let ast::Expr::FieldExpr(field_expr) = expr_type.expr()?
                 && let Some(field) = field_expr.field()
                 && let Some(ast::Expr::NameRef(schema_name_ref)) = field_expr.base()
             {