Merge pull request #19657 from ChayimFriedman2/better-offset-of

feat: Better support `offset_of!()`

crates/hir/src/semantics.rs

@@ -1622,6 +1622,13 @@ impl<'db> SemanticsImpl<'db> {
         self.analyze(name.syntax())?.resolve_use_type_arg(name)
     }
 
+    pub fn resolve_offset_of_field(
+        &self,
+        name_ref: &ast::NameRef,
+    ) -> Option<(Either<Variant, Field>, GenericSubstitution)> {
+        self.analyze_no_infer(name_ref.syntax())?.resolve_offset_of_field(self.db, name_ref)
+    }
+
     pub fn resolve_mod_path(
         &self,
         scope: &SyntaxNode,

crates/hir/src/source_analyzer.rs

@@ -14,7 +14,7 @@ use crate::{
 };
 use either::Either;
 use hir_def::{
-    AssocItemId, CallableDefId, ConstId, DefWithBodyId, FieldId, FunctionId, GenericDefId,
+    AdtId, AssocItemId, CallableDefId, ConstId, DefWithBodyId, FieldId, FunctionId, GenericDefId,
     ItemContainerId, LocalFieldId, Lookup, ModuleDefId, StructId, TraitId, VariantId,
     expr_store::{
         Body, BodySourceMap, ExpressionStore, ExpressionStoreSourceMap, HygieneId,
@@ -34,8 +34,8 @@ use hir_expand::{
     name::{AsName, Name},
 };
 use hir_ty::{
-    Adjustment, InferenceResult, Interner, Substitution, TraitEnvironment, Ty, TyExt, TyKind,
+    Adjustment, AliasTy, InferenceResult, Interner, ProjectionTy, Substitution, TraitEnvironment,
-    TyLoweringContext,
+    Ty, TyExt, TyKind, TyLoweringContext,
     diagnostics::{
         InsideUnsafeBlock, record_literal_missing_fields, record_pattern_missing_fields,
         unsafe_operations,
@@ -47,6 +47,7 @@ use hir_ty::{
 use intern::sym;
 use itertools::Itertools;
 use smallvec::SmallVec;
+use stdx::never;
 use syntax::{
     SyntaxKind, SyntaxNode, TextRange, TextSize,
     ast::{self, AstNode, RangeItem, RangeOp},
@@ -788,6 +789,78 @@ impl SourceAnalyzer {
             .map(crate::TypeParam::from)
     }
 
+    pub(crate) fn resolve_offset_of_field(
+        &self,
+        db: &dyn HirDatabase,
+        name_ref: &ast::NameRef,
+    ) -> Option<(Either<crate::Variant, crate::Field>, GenericSubstitution)> {
+        let offset_of_expr = ast::OffsetOfExpr::cast(name_ref.syntax().parent()?)?;
+        let container = offset_of_expr.ty()?;
+        let container = self.type_of_type(db, &container)?;
+
+        let trait_env = container.env;
+        let mut container = Either::Right(container.ty);
+        for field_name in offset_of_expr.fields() {
+            if let Some(
+                TyKind::Alias(AliasTy::Projection(ProjectionTy { associated_ty_id, substitution }))
+                | TyKind::AssociatedType(associated_ty_id, substitution),
+            ) = container.as_ref().right().map(|it| it.kind(Interner))
+            {
+                let projection = ProjectionTy {
+                    associated_ty_id: *associated_ty_id,
+                    substitution: substitution.clone(),
+                };
+                container = Either::Right(db.normalize_projection(projection, trait_env.clone()));
+            }
+            let handle_variants = |variant, subst: &Substitution, container: &mut _| {
+                let fields = db.variant_fields(variant);
+                let field = fields.field(&field_name.as_name())?;
+                let field_types = db.field_types(variant);
+                *container = Either::Right(field_types[field].clone().substitute(Interner, subst));
+                let generic_def = match variant {
+                    VariantId::EnumVariantId(it) => it.loc(db).parent.into(),
+                    VariantId::StructId(it) => it.into(),
+                    VariantId::UnionId(it) => it.into(),
+                };
+                Some((
+                    Either::Right(Field { parent: variant.into(), id: field }),
+                    generic_def,
+                    subst.clone(),
+                ))
+            };
+            let temp_ty = TyKind::Error.intern(Interner);
+            let (field_def, generic_def, subst) =
+                match std::mem::replace(&mut container, Either::Right(temp_ty.clone())) {
+                    Either::Left((variant_id, subst)) => {
+                        handle_variants(VariantId::from(variant_id), &subst, &mut container)?
+                    }
+                    Either::Right(container_ty) => match container_ty.kind(Interner) {
+                        TyKind::Adt(adt_id, subst) => match adt_id.0 {
+                            AdtId::StructId(id) => {
+                                handle_variants(id.into(), subst, &mut container)?
+                            }
+                            AdtId::UnionId(id) => {
+                                handle_variants(id.into(), subst, &mut container)?
+                            }
+                            AdtId::EnumId(id) => {
+                                let variants = db.enum_variants(id);
+                                let variant = variants.variant(&field_name.as_name())?;
+                                container = Either::Left((variant, subst.clone()));
+                                (Either::Left(Variant { id: variant }), id.into(), subst.clone())
+                            }
+                        },
+                        _ => return None,
+                    },
+                };
+
+            if field_name.syntax().text_range() == name_ref.syntax().text_range() {
+                return Some((field_def, GenericSubstitution::new(generic_def, subst, trait_env)));
+            }
+        }
+        never!("the `NameRef` is a child of the `OffsetOfExpr`, we should've visited it");
+        None
+    }
+
     pub(crate) fn resolve_path(
         &self,
         db: &dyn HirDatabase,

crates/ide-db/src/defs.rs

@@ -839,6 +839,14 @@ impl NameRefClass {
                 ast::AsmRegSpec(_) => {
                     Some(NameRefClass::Definition(Definition::InlineAsmRegOrRegClass(()), None))
                 },
+                ast::OffsetOfExpr(_) => {
+                    let (def, subst) = sema.resolve_offset_of_field(name_ref)?;
+                    let def = match def {
+                        Either::Left(variant) => Definition::Variant(variant),
+                        Either::Right(field) => Definition::Field(field),
+                    };
+                    Some(NameRefClass::Definition(def, Some(subst)))
+                },
                 _ => None
             }
         }

crates/ide/src/goto_definition.rs

@@ -3472,6 +3472,76 @@ fn main() {
 
         fn f1<const N: buz$0>() {}
     }
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn offset_of() {
+        check(
+            r#"
+//- minicore: offset_of
+struct Foo {
+    field: i32,
+ // ^^^^^
+}
+
+fn foo() {
+    let _ = core::mem::offset_of!(Foo, fiel$0d);
+}
+        "#,
+        );
+
+        check(
+            r#"
+//- minicore: offset_of
+struct Bar(Foo);
+struct Foo {
+    field: i32,
+ // ^^^^^
+}
+
+fn foo() {
+    let _ = core::mem::offset_of!(Bar, 0.fiel$0d);
+}
+        "#,
+        );
+
+        check(
+            r#"
+//- minicore: offset_of
+struct Bar(Baz);
+enum Baz {
+    Abc(Foo),
+    None,
+}
+struct Foo {
+    field: i32,
+ // ^^^^^
+}
+
+fn foo() {
+    let _ = core::mem::offset_of!(Bar, 0.Abc.0.fiel$0d);
+}
+        "#,
+        );
+
+        check(
+            r#"
+//- minicore: offset_of
+struct Bar(Baz);
+enum Baz {
+    Abc(Foo),
+ // ^^^
+    None,
+}
+struct Foo {
+    field: i32,
+}
+
+fn foo() {
+    let _ = core::mem::offset_of!(Bar, 0.Ab$0c.0.field);
 }
         "#,
         );

crates/parser/src/grammar/expressions/atom.rs

@@ -258,6 +258,15 @@ fn builtin_expr(p: &mut Parser<'_>) -> Option<CompletedMarker> {
         p.expect(T!['(']);
         type_(p);
         p.expect(T![,]);
+        // Due to our incomplete handling of macro groups, especially
+        // those with empty delimiters, we wrap `expr` fragments in
+        // parentheses sometimes. Since `offset_of` is a macro, and takes
+        // `expr`, the field names could be wrapped in parentheses.
+        let wrapped_in_parens = p.eat(T!['(']);
+        // test offset_of_parens
+        // fn foo() {
+        //     builtin#offset_of(Foo, (bar.baz.0));
+        // }
         while !p.at(EOF) && !p.at(T![')']) {
             name_ref_mod_path_or_index(p);
             if !p.at(T![')']) {
@@ -265,6 +274,9 @@ fn builtin_expr(p: &mut Parser<'_>) -> Option<CompletedMarker> {
             }
         }
         p.expect(T![')']);
+        if wrapped_in_parens {
+            p.expect(T![')']);
+        }
         Some(m.complete(p, OFFSET_OF_EXPR))
     } else if p.at_contextual_kw(T![format_args]) {
         p.bump_remap(T![format_args]);

crates/parser/test_data/generated/runner.rs

@@ -422,6 +422,10 @@ mod ok {
         run_and_expect_no_errors("test_data/parser/inline/ok/nocontentexpr_after_item.rs");
     }
     #[test]
+    fn offset_of_parens() {
+        run_and_expect_no_errors("test_data/parser/inline/ok/offset_of_parens.rs");
+    }
+    #[test]
     fn or_pattern() { run_and_expect_no_errors("test_data/parser/inline/ok/or_pattern.rs"); }
     #[test]
     fn param_list() { run_and_expect_no_errors("test_data/parser/inline/ok/param_list.rs"); }

crates/parser/test_data/parser/inline/ok/offset_of_parens.rast

@@ -0,0 +1,42 @@
+SOURCE_FILE
+  FN
+    FN_KW "fn"
+    WHITESPACE " "
+    NAME
+      IDENT "foo"
+    PARAM_LIST
+      L_PAREN "("
+      R_PAREN ")"
+    WHITESPACE " "
+    BLOCK_EXPR
+      STMT_LIST
+        L_CURLY "{"
+        WHITESPACE "\n    "
+        EXPR_STMT
+          OFFSET_OF_EXPR
+            BUILTIN_KW "builtin"
+            POUND "#"
+            OFFSET_OF_KW "offset_of"
+            L_PAREN "("
+            PATH_TYPE
+              PATH
+                PATH_SEGMENT
+                  NAME_REF
+                    IDENT "Foo"
+            COMMA ","
+            WHITESPACE " "
+            L_PAREN "("
+            NAME_REF
+              IDENT "bar"
+            DOT "."
+            NAME_REF
+              IDENT "baz"
+            DOT "."
+            NAME_REF
+              INT_NUMBER "0"
+            R_PAREN ")"
+            R_PAREN ")"
+          SEMICOLON ";"
+        WHITESPACE "\n"
+        R_CURLY "}"
+  WHITESPACE "\n"

crates/parser/test_data/parser/inline/ok/offset_of_parens.rs

@@ -0,0 +1,3 @@
+fn foo() {
+    builtin#offset_of(Foo, (bar.baz.0));
+}

crates/test-utils/src/minicore.rs

@@ -70,6 +70,7 @@
 //!     unimplemented: panic
 //!     column:
 //!     addr_of:
+//!     offset_of:
 
 #![rustc_coherence_is_core]
 
@@ -414,6 +415,13 @@ pub mod mem {
     use crate::marker::DiscriminantKind;
     pub struct Discriminant<T>(<T as DiscriminantKind>::Discriminant);
     // endregion:discriminant
+
+    // region:offset_of
+    pub macro offset_of($Container:ty, $($fields:expr)+ $(,)?) {
+        // The `{}` is for better error messages
+        {builtin # offset_of($Container, $($fields)+)}
+    }
+    // endregion:offset_of
 }
 
 pub mod ptr {