From 19b7a7da12c15037c174960948708ac5d131c9ba Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 19 May 2023 08:34:53 -0700
Subject: [PATCH] Fix regression on concrete playback inplace (#2454)

This fixes a regression from #2439. The compiler should store the location of the function body instead of the declaration. Storing the correct location fixes how concrete playback stores the generated unit test.
---
 kani-compiler/src/kani_middle/metadata.rs |  8 +++++---
 kani-compiler/src/kani_middle/mod.rs      | 10 +---------
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/kani-compiler/src/kani_middle/metadata.rs b/kani-compiler/src/kani_middle/metadata.rs
index de3667c6e6a8..e293c95e2cd8 100644
--- a/kani-compiler/src/kani_middle/metadata.rs
+++ b/kani-compiler/src/kani_middle/metadata.rs
@@ -8,7 +8,7 @@ use std::path::Path;
 use crate::kani_middle::attributes::test_harness_name;
 use kani_metadata::{ArtifactType, HarnessAttributes, HarnessMetadata};
 use rustc_hir::def_id::DefId;
-use rustc_middle::ty::{Instance, TyCtxt};
+use rustc_middle::ty::{Instance, InstanceDef, TyCtxt, WithOptConstParam};
 
 use super::{attributes::extract_harness_attributes, SourceLocation};
 
@@ -24,7 +24,8 @@ pub fn gen_proof_metadata(tcx: TyCtxt, def_id: DefId, base_name: &Path) -> Harne
         tcx.symbol_name(Instance::mono(tcx, def_id)).to_string()
     };
 
-    let loc = SourceLocation::def_id_loc(tcx, def_id);
+    let body = tcx.instance_mir(InstanceDef::Item(WithOptConstParam::unknown(def_id)));
+    let loc = SourceLocation::new(tcx, &body.span);
     let file_stem = format!("{}_{mangled_name}", base_name.file_stem().unwrap().to_str().unwrap());
     let model_file = base_name.with_file_name(file_stem).with_extension(ArtifactType::SymTabGoto);
 
@@ -51,7 +52,8 @@ pub fn gen_test_metadata<'tcx>(
 ) -> HarnessMetadata {
     let pretty_name = test_harness_name(tcx, test_desc);
     let mangled_name = tcx.symbol_name(test_fn).to_string();
-    let loc = SourceLocation::def_id_loc(tcx, test_desc);
+    let body = tcx.instance_mir(InstanceDef::Item(WithOptConstParam::unknown(test_desc)));
+    let loc = SourceLocation::new(tcx, &body.span);
     let file_stem = format!("{}_{mangled_name}", base_name.file_stem().unwrap().to_str().unwrap());
     let model_file = base_name.with_file_name(file_stem).with_extension(ArtifactType::SymTabGoto);
 
diff --git a/kani-compiler/src/kani_middle/mod.rs b/kani-compiler/src/kani_middle/mod.rs
index 358ec2149f67..fbe27bd8102b 100644
--- a/kani-compiler/src/kani_middle/mod.rs
+++ b/kani-compiler/src/kani_middle/mod.rs
@@ -6,10 +6,7 @@
 use std::collections::HashSet;
 
 use kani_queries::{QueryDb, UserInput};
-use rustc_hir::{
-    def::DefKind,
-    def_id::{DefId, LOCAL_CRATE},
-};
+use rustc_hir::{def::DefKind, def_id::LOCAL_CRATE};
 use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::span_bug;
 use rustc_middle::ty::layout::{
@@ -103,11 +100,6 @@ impl SourceLocation {
         };
         SourceLocation { filename, start_line, start_col, end_line, end_col }
     }
-
-    pub fn def_id_loc(tcx: TyCtxt, def_id: DefId) -> Self {
-        let span = tcx.def_span(def_id);
-        Self::new(tcx, &span)
-    }
 }
 
 /// Get the FnAbi of a given instance with no extra variadic arguments.