fix: propagate the DefId for the owner function of formula_fn and

extern_spec_fn

Author: coeff-aij

src/analyze.rs

@@ -412,6 +412,7 @@ impl<'tcx> Analyzer<'tcx> {
         &self,
         local_def_id: LocalDefId,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        owner_fn_id: DefId,
     ) -> Option<annot_fn::FormulaFn<'tcx>> {
         let deferred_formula_fn = self.formula_fns.get(&local_def_id)?;
 
@@ -425,6 +426,7 @@ impl<'tcx> Analyzer<'tcx> {
             local_def_id,
             self.type_params.clone(),
             self.system.clone(),
+            owner_fn_id,
         )
         .with_generic_args(generic_args)
         .with_def_id_cache(self.def_ids());
@@ -548,7 +550,7 @@ impl<'tcx> Analyzer<'tcx> {
         &mut self,
         local_def_id: LocalDefId,
         bb: BasicBlock,
-        owner_fn_id: DefId
+        owner_fn_id: DefId,
     ) -> basic_block::Analyzer<'tcx, '_> {
         basic_block::Analyzer::new(self, local_def_id, bb, owner_fn_id)
     }
@@ -654,6 +656,7 @@ impl<'tcx> Analyzer<'tcx> {
         resolver: T,
         self_type_name: Option<String>,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        owner_fn_id: DefId,
     ) -> Option<AnnotFormula<T::Output>>
     where
         T: Resolver<Output = rty::FunctionParamIdx>,
@@ -684,7 +687,9 @@ impl<'tcx> Analyzer<'tcx> {
             if require_annot.is_some() {
                 unimplemented!();
             }
-            let Some(formula_fn) = self.formula_fn_with_args(formula_def_id, generic_args) else {
+            let Some(formula_fn) =
+                self.formula_fn_with_args(formula_def_id, generic_args, owner_fn_id)
+            else {
                 panic!(
                     "require annotation {:?} is not a formula function",
                     formula_def_id
@@ -703,6 +708,7 @@ impl<'tcx> Analyzer<'tcx> {
         resolver: T,
         self_type_name: Option<String>,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        owner_fn_id: DefId,
     ) -> Option<AnnotFormula<T::Output>>
     where
         T: Resolver<Output = rty::RefinedTypeVar<rty::FunctionParamIdx>>,
@@ -734,7 +740,9 @@ impl<'tcx> Analyzer<'tcx> {
             if ensure_annot.is_some() {
                 unimplemented!();
             }
-            let Some(formula_fn) = self.formula_fn_with_args(formula_def_id, generic_args) else {
+            let Some(formula_fn) =
+                self.formula_fn_with_args(formula_def_id, generic_args, owner_fn_id)
+            else {
                 panic!(
                     "ensure annotation {:?} is not a formula function",
                     formula_def_id

src/analyze/annot_fn.rs

@@ -3,7 +3,10 @@ use std::collections::HashMap;
 use std::rc::Rc;
 
 use pretty::{termcolor, Pretty};
-use rustc_hir::{def_id::LocalDefId, HirId};
+use rustc_hir::{
+    def_id::{DefId, LocalDefId},
+    HirId,
+};
 use rustc_index::IndexVec;
 use rustc_middle::ty::{self as mir_ty, TyCtxt, TypeFoldable};
 
@@ -150,6 +153,7 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         local_def_id: LocalDefId,
         type_params: Rc<RefCell<TypeParamMap>>,
         system: Rc<RefCell<System>>,
+        owner_fn_id: DefId,
     ) -> Self {
         let body = tcx.hir_body_owned_by(local_def_id);
         let generic_args = tcx.mk_args(&[]);
@@ -158,7 +162,7 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         let type_builder = TypeBuilder::new(
             tcx,
             def_ids.clone(),
-            local_def_id.to_def_id(),
+            owner_fn_id,
             type_params.clone(),
             system.clone(),
         );

src/analyze/basic_block.rs

@@ -649,7 +649,10 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             );
         }
         tracing::info!(?def_id, ?resolved_def_id, ?resolved_args, "resolved");
-        let Some(def_ty) = self.ctx.def_ty_with_args(resolved_def_id, resolved_args, caller_def_id) else {
+        let Some(def_ty) = self
+            .ctx
+            .def_ty_with_args(resolved_def_id, resolved_args, caller_def_id)
+        else {
             panic!(
                 "unknown def (resolved): {:?}, args: {:?}",
                 resolved_def_id, resolved_args

src/analyze/crate_.rs

@@ -111,14 +111,9 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             }
         }
 
-        let target_def_id = if analyzer.is_annotated_as_extern_spec_fn() {
-            analyzer.extern_spec_fn_target_def_id()
-        } else {
-            local_def_id.to_def_id()
-        };
-
+        let owner_fn_id = analyzer.owner_fn_id;
         let expected = analyzer.expected_ty();
-        self.ctx.register_def(target_def_id, expected);
+        self.ctx.register_def(owner_fn_id, expected);
     }
 
     fn analyze_local_defs(&mut self) {
@@ -140,7 +135,6 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             // check polymorphic function def by replacing type params with some opaque type
             // (and this is no-op if the function is mono)
             let expected = expected.clone();
-            tracing::debug!("expected type of {:?} is {:#?}", local_def_id, expected);
             let generic_args = self.placeholder_generic_args(*local_def_id);
             self.ctx
                 .local_def_analyzer(*local_def_id)

src/analyze/local_def.rs

@@ -37,6 +37,82 @@ fn stmt_str_literal(stmt: &rustc_hir::Stmt) -> Option<String> {
     }
 }
 
+fn is_annotated_as_extern_spec_fn_impl(tcx: &TyCtxt, local_def_id: &LocalDefId) -> bool {
+    tcx.get_attrs_by_path(
+        local_def_id.to_def_id(),
+        &analyze::annot::extern_spec_fn_path(),
+    )
+    .next()
+    .is_some()
+}
+
+/// Extract the target DefId from `#[thrust::extern_spec_fn]` function.
+///
+/// The target is identified as the tail call expression (last expression without
+/// semicolon) in the function body block.
+fn extern_spec_fn_target_def_id_impl<'tcx>(
+    tcx: &TyCtxt<'tcx>,
+    local_def_id: &LocalDefId,
+    mir_body: &Body<'tcx>,
+) -> DefId {
+    let hir_node = tcx.hir_node_by_def_id(*local_def_id);
+    let hir_body_id = match hir_node {
+        rustc_hir::Node::Item(item) => {
+            let rustc_hir::ItemKind::Fn { body: body_id, .. } = item.kind else {
+                panic!("extern_spec_fn must be a function");
+            };
+            body_id
+        }
+        rustc_hir::Node::ImplItem(impl_item) => {
+            let rustc_hir::ImplItemKind::Fn(_, body_id) = impl_item.kind else {
+                panic!("extern_spec_fn must be a function");
+            };
+            body_id
+        }
+        rustc_hir::Node::TraitItem(trait_item) => {
+            let rustc_hir::TraitItemKind::Fn(_, rustc_hir::TraitFn::Provided(body_id)) =
+                trait_item.kind
+            else {
+                panic!("extern_spec_fn must be a function with a body");
+            };
+            body_id
+        }
+        _ => panic!("extern_spec_fn must be a function item or impl item"),
+    };
+
+    let hir_body = tcx.hir_body(hir_body_id);
+
+    // The body is a block; the tail expression is the function call to the target.
+    let rustc_hir::ExprKind::Block(block, _) = &hir_body.value.kind else {
+        panic!("extern_spec_fn body must be a block");
+    };
+    let tail_expr = block
+        .expr
+        .expect("extern_spec_fn block must end with a tail call expression");
+
+    let rustc_hir::ExprKind::Call(func_expr, _) = &tail_expr.kind else {
+        panic!("extern_spec_fn tail expression must be a function call");
+    };
+    let rustc_hir::ExprKind::Path(qpath) = &func_expr.kind else {
+        panic!("extern_spec_fn call must be a path expression");
+    };
+
+    let typeck_result = tcx.typeck(local_def_id);
+    let hir_id = func_expr.hir_id;
+    let rustc_hir::def::Res::Def(_, def_id) = typeck_result.qpath_res(qpath, hir_id) else {
+        panic!("extern_spec_fn call must resolve to a definition");
+    };
+
+    let args = typeck_result.node_args(hir_id);
+    let typing_env = mir_body.typing_env(*tcx);
+    let instance = mir_ty::Instance::try_resolve(*tcx, typing_env, def_id, args).unwrap();
+    if let Some(instance) = instance {
+        instance.def_id()
+    } else {
+        def_id
+    }
+}
+
 /// An implementation of the typing of local definitions.
 ///
 /// The current implementation only applies to function definitions. The entry point is
@@ -46,6 +122,7 @@ pub struct Analyzer<'tcx, 'ctx> {
     tcx: TyCtxt<'tcx>,
 
     local_def_id: LocalDefId,
+    pub owner_fn_id: DefId,
 
     body: Body<'tcx>,
     /// to substitute HIR types during translation in [`crate::analyze::annot_fn`]
@@ -201,13 +278,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     }
 
     pub fn is_annotated_as_extern_spec_fn(&self) -> bool {
-        self.tcx
-            .get_attrs_by_path(
-                self.local_def_id.to_def_id(),
-                &analyze::annot::extern_spec_fn_path(),
-            )
-            .next()
-            .is_some()
+        is_annotated_as_extern_spec_fn_impl(&self.tcx, &self.local_def_id)
     }
 
     pub fn is_annotated_as_predicate(&self) -> bool {
@@ -318,7 +389,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             .associated_item(self.local_def_id.to_def_id())
             .trait_item_def_id
             .unwrap();
-        self.ctx.def_ty_with_args(trait_item_did, trait_ref.args, trait_ref.def_id)
+        self.ctx
+            .def_ty_with_args(trait_item_did, trait_ref.args, trait_ref.def_id)
     }
 
     // Note that we do not expect predicate variables to be generated here
@@ -358,13 +430,15 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             &param_resolver,
             self_type_name.clone(),
             self.generic_args,
+            self.owner_fn_id,
         );
 
         let mut ensure_annot = self.ctx.extract_ensure_annot(
             self.local_def_id,
             &result_param_resolver,
             self_type_name.clone(),
             self.generic_args,
+            self.owner_fn_id,
         );
 
         if let Some(trait_item_id) = self.local_trait_item_id() {
@@ -374,12 +448,14 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 &param_resolver,
                 self_type_name.clone(),
                 self.generic_args,
+                self.owner_fn_id,
             );
             let trait_ensure_annot = self.ctx.extract_ensure_annot(
                 trait_item_id,
                 &result_param_resolver,
                 self_type_name.clone(),
                 self.generic_args,
+                self.owner_fn_id,
             );
 
             assert!(require_annot.is_none() || trait_require_annot.is_none());
@@ -447,62 +523,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     /// The target is identified as the tail call expression (last expression without
     /// semicolon) in the function body block.
     pub fn extern_spec_fn_target_def_id(&self) -> DefId {
-        let node = self.tcx.hir_node_by_def_id(self.local_def_id);
-        let body_id = match node {
-            rustc_hir::Node::Item(item) => {
-                let rustc_hir::ItemKind::Fn { body: body_id, .. } = item.kind else {
-                    panic!("extern_spec_fn must be a function");
-                };
-                body_id
-            }
-            rustc_hir::Node::ImplItem(impl_item) => {
-                let rustc_hir::ImplItemKind::Fn(_, body_id) = impl_item.kind else {
-                    panic!("extern_spec_fn must be a function");
-                };
-                body_id
-            }
-            rustc_hir::Node::TraitItem(trait_item) => {
-                let rustc_hir::TraitItemKind::Fn(_, rustc_hir::TraitFn::Provided(body_id)) =
-                    trait_item.kind
-                else {
-                    panic!("extern_spec_fn must be a function with a body");
-                };
-                body_id
-            }
-            _ => panic!("extern_spec_fn must be a function item or impl item"),
-        };
-
-        let body = self.tcx.hir_body(body_id);
-
-        // The body is a block; the tail expression is the function call to the target.
-        let rustc_hir::ExprKind::Block(block, _) = &body.value.kind else {
-            panic!("extern_spec_fn body must be a block");
-        };
-        let tail_expr = block
-            .expr
-            .expect("extern_spec_fn block must end with a tail call expression");
-
-        let rustc_hir::ExprKind::Call(func_expr, _) = &tail_expr.kind else {
-            panic!("extern_spec_fn tail expression must be a function call");
-        };
-        let rustc_hir::ExprKind::Path(qpath) = &func_expr.kind else {
-            panic!("extern_spec_fn call must be a path expression");
-        };
-
-        let typeck_result = self.tcx.typeck(self.local_def_id);
-        let hir_id = func_expr.hir_id;
-        let rustc_hir::def::Res::Def(_, def_id) = typeck_result.qpath_res(qpath, hir_id) else {
-            panic!("extern_spec_fn call must resolve to a definition");
-        };
-
-        let args = typeck_result.node_args(hir_id);
-        let typing_env = self.body.typing_env(self.tcx);
-        let instance = mir_ty::Instance::try_resolve(self.tcx, typing_env, def_id, args).unwrap();
-        if let Some(instance) = instance {
-            instance.def_id()
-        } else {
-            def_id
-        }
+        extern_spec_fn_target_def_id_impl(&self.tcx, &self.local_def_id, &self.body)
     }
 
     fn is_mut_param(&self, param_idx: rty::FunctionParamIdx) -> bool {
@@ -973,12 +994,18 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         let tcx = ctx.tcx;
         let body = tcx.optimized_mir(local_def_id.to_def_id()).clone();
         let drop_points = Default::default();
-        let type_builder = ctx.type_builder(ctx.def_ids(), local_def_id.to_def_id());
+        let owner_fn_id = if is_annotated_as_extern_spec_fn_impl(&tcx, &local_def_id) {
+            extern_spec_fn_target_def_id_impl(&tcx, &local_def_id, &body)
+        } else {
+            local_def_id.to_def_id()
+        };
+        let type_builder = ctx.type_builder(ctx.def_ids(), owner_fn_id);
         let generic_args = tcx.mk_args(&[]);
         Self {
             ctx,
             tcx,
             local_def_id,
+            owner_fn_id,
             body,
             generic_args,
             drop_points,