Implement offset_of in hir-def and hir-ty

Author: Veykril

crates/hir-def/src/body/lower.rs

@@ -31,7 +31,7 @@ use crate::{
     hir::{
         dummy_expr_id, Array, Binding, BindingAnnotation, BindingId, BindingProblems, CaptureBy,
         ClosureKind, Expr, ExprId, Label, LabelId, Literal, LiteralOrConst, MatchArm, Movability,
-        Pat, PatId, RecordFieldPat, RecordLitField, Statement,
+        OffsetOf, Pat, PatId, RecordFieldPat, RecordLitField, Statement,
     },
     item_scope::BuiltinShadowMode,
     lang_item::LangItem,
@@ -649,7 +649,11 @@ impl ExprCollector<'_> {
             }
             ast::Expr::UnderscoreExpr(_) => self.alloc_expr(Expr::Underscore, syntax_ptr),
             ast::Expr::AsmExpr(_) => self.missing_expr(),
-            ast::Expr::OffsetOfExpr(_) => self.missing_expr(),
+            ast::Expr::OffsetOfExpr(e) => {
+                let container = Interned::new(TypeRef::from_ast_opt(&self.ctx(), e.ty()));
+                let fields = e.fields().map(|it| it.as_name()).collect();
+                self.alloc_expr(Expr::OffsetOf(OffsetOf { container, fields }), syntax_ptr)
+            }
             ast::Expr::FormatArgsExpr(_) => self.missing_expr(),
         })
     }

crates/hir-def/src/body/pretty.rs

@@ -3,6 +3,7 @@
 use std::fmt::{self, Write};
 
 use hir_expand::db::ExpandDatabase;
+use itertools::Itertools;
 use syntax::ast::HasName;
 
 use crate::{
@@ -154,6 +155,15 @@ impl Printer<'_> {
         match expr {
             Expr::Missing => w!(self, "�"),
             Expr::Underscore => w!(self, "_"),
+            Expr::OffsetOf(offset_of) => {
+                w!(self, "builtin#offset_of!(");
+                self.print_type_ref(&offset_of.container);
+                w!(
+                    self,
+                    ", {})",
+                    offset_of.fields.iter().format_with(".", |field, f| f(&field.display(self.db)))
+                );
+            }
             Expr::Path(path) => self.print_path(path),
             Expr::If { condition, then_branch, else_branch } => {
                 w!(self, "if ");

crates/hir-def/src/hir.rs

@@ -281,6 +281,13 @@ pub enum Expr {
     Array(Array),
     Literal(Literal),
     Underscore,
+    OffsetOf(OffsetOf),
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct OffsetOf {
+    pub container: Interned<TypeRef>,
+    pub fields: Box<[Name]>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -341,7 +348,7 @@ impl Expr {
     pub fn walk_child_exprs(&self, mut f: impl FnMut(ExprId)) {
         match self {
             Expr::Missing => {}
-            Expr::Path(_) => {}
+            Expr::Path(_) | Expr::OffsetOf(_) => {}
             Expr::If { condition, then_branch, else_branch } => {
                 f(*condition);
                 f(*then_branch);

crates/hir-ty/src/infer/closure.rs

@@ -452,6 +452,7 @@ impl InferenceContext<'_> {
 
     fn walk_expr_without_adjust(&mut self, tgt_expr: ExprId) {
         match &self.body[tgt_expr] {
+            Expr::OffsetOf(_) => (),
             Expr::If { condition, then_branch, else_branch } => {
                 self.consume_expr(*condition);
                 self.consume_expr(*then_branch);

crates/hir-ty/src/infer/expr.rs

@@ -843,6 +843,7 @@ impl InferenceContext<'_> {
                 });
                 expected
             }
+            Expr::OffsetOf(_) => TyKind::Scalar(Scalar::Uint(UintTy::Usize)).intern(Interner),
         };
         // use a new type variable if we got unknown here
         let ty = self.insert_type_vars_shallow(ty);

crates/hir-ty/src/infer/mutability.rs

@@ -35,6 +35,7 @@ impl InferenceContext<'_> {
     fn infer_mut_expr_without_adjust(&mut self, tgt_expr: ExprId, mutability: Mutability) {
         match &self.body[tgt_expr] {
             Expr::Missing => (),
+            Expr::OffsetOf(_) => (),
             &Expr::If { condition, then_branch, else_branch } => {
                 self.infer_mut_expr(condition, Mutability::Not);
                 self.infer_mut_expr(then_branch, Mutability::Not);

crates/hir-ty/src/mir/lower.rs

@@ -370,6 +370,7 @@ impl<'ctx> MirLowerCtx<'ctx> {
         mut current: BasicBlockId,
     ) -> Result<Option<BasicBlockId>> {
         match &self.body.exprs[expr_id] {
+            Expr::OffsetOf(_) => not_supported!("builtin#offset_of"),
             Expr::Missing => {
                 if let DefWithBodyId::FunctionId(f) = self.owner {
                     let assoc = f.lookup(self.db.upcast());

crates/hir-ty/src/tests/simple.rs

@@ -3649,3 +3649,15 @@ fn main() {
 "#,
     );
 }
+
+#[test]
+fn offset_of() {
+    check_types(
+        r#"
+fn main() {
+    builtin#offset_of((,), 0);
+ // ^^^^^^^^^^^^^^^^^^^^^^^^^ usize
+}
+"#,
+    );
+}

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

@@ -220,24 +220,39 @@ fn tuple_expr(p: &mut Parser<'_>) -> CompletedMarker {
 // fn foo() {
 //     builtin#asm(0);
 //     builtin#format_args(0);
-//     builtin#builtin(0);
+//     builtin#offset_of(Foo, bar.baz.0);
 // }
 fn builtin_expr(p: &mut Parser<'_>) -> Option<CompletedMarker> {
     let m = p.start();
     p.bump_remap(T![builtin]);
     p.bump(T![#]);
     if p.at_contextual_kw(T![offset_of]) {
+        p.bump_remap(T![offset_of]);
         p.expect(T!['(']);
         type_(p);
-        p.bump(T![,]);
+        p.expect(T![,]);
+        while !p.at(EOF) && !p.at(T![')']) {
+            if p.at(IDENT) || p.at(INT_NUMBER) {
+                name_ref_or_index(p);
+            // } else if p.at(FLOAT_NUMBER) {
+            // FIXME: needs float hack
+            } else {
+                p.err_and_bump("expected field name or number");
+            }
+            if !p.at(T![')']) {
+                p.expect(T![.]);
+            }
+        }
         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]);
         p.expect(T!['(']);
         expr(p);
         p.expect(T![')']);
         Some(m.complete(p, FORMAT_ARGS_EXPR))
     } else if p.at_contextual_kw(T![asm]) {
+        p.bump_remap(T![asm]);
         p.expect(T!['(']);
         expr(p);
         p.expect(T![')']);