From 99a4f2e98334a332a25b42221a79c12b67a630d8 Mon Sep 17 00:00:00 2001
From: hkalbasi <hamidrezakalbasi@protonmail.com>
Date: Thu, 22 Jun 2023 22:24:21 +0330
Subject: [PATCH] Support manual impl of fn traits in mir interpreter

---
 crates/hir-ty/src/consteval/tests.rs          | 24 ++++++++++++
 .../hir-ty/src/consteval/tests/intrinsics.rs  |  1 +
 crates/hir-ty/src/mir/eval.rs                 | 39 +++++++++++++++++--
 crates/hir-ty/src/mir/eval/shim.rs            | 18 ++++++++-
 4 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/crates/hir-ty/src/consteval/tests.rs b/crates/hir-ty/src/consteval/tests.rs
index 1e241500c2c..b487d7d049b 100644
--- a/crates/hir-ty/src/consteval/tests.rs
+++ b/crates/hir-ty/src/consteval/tests.rs
@@ -1550,6 +1550,30 @@ fn closures() {
     );
 }
 
+#[test]
+fn manual_fn_trait_impl() {
+    check_number(
+        r#"
+//- minicore: fn, copy
+struct S(i32);
+
+impl FnOnce<(i32, i32)> for S {
+    type Output = i32;
+
+    extern "rust-call" fn call_once(self, arg: (i32, i32)) -> i32 {
+        arg.0 + arg.1 + self.0
+    }
+}
+
+const GOAL: i32 = {
+    let s = S(1);
+    s(2, 3)
+};
+"#,
+        6,
+    );
+}
+
 #[test]
 fn closure_and_impl_fn() {
     check_number(
diff --git a/crates/hir-ty/src/consteval/tests/intrinsics.rs b/crates/hir-ty/src/consteval/tests/intrinsics.rs
index e05d824dbac..24a463d714c 100644
--- a/crates/hir-ty/src/consteval/tests/intrinsics.rs
+++ b/crates/hir-ty/src/consteval/tests/intrinsics.rs
@@ -32,6 +32,7 @@ fn transmute() {
 fn const_eval_select() {
     check_number(
         r#"
+        //- minicore: fn
         extern "rust-intrinsic" {
             pub fn const_eval_select<ARG, F, G, RET>(arg: ARG, called_in_const: F, called_at_rt: G) -> RET
             where
diff --git a/crates/hir-ty/src/mir/eval.rs b/crates/hir-ty/src/mir/eval.rs
index 28b6066c191..0c7669cd54d 100644
--- a/crates/hir-ty/src/mir/eval.rs
+++ b/crates/hir-ty/src/mir/eval.rs
@@ -1799,7 +1799,7 @@ impl Evaluator<'_> {
         match def {
             CallableDefId::FunctionId(def) => {
                 if let Some(_) = self.detect_fn_trait(def) {
-                    self.exec_fn_trait(&args, destination, locals, span)?;
+                    self.exec_fn_trait(def, args, generic_args, locals, destination, span)?;
                     return Ok(());
                 }
                 self.exec_fn_with_args(def, args, generic_args, locals, destination, span)?;
@@ -1921,9 +1921,11 @@ impl Evaluator<'_> {
 
     fn exec_fn_trait(
         &mut self,
+        def: FunctionId,
         args: &[IntervalAndTy],
-        destination: Interval,
+        generic_args: Substitution,
         locals: &Locals<'_>,
+        destination: Interval,
         span: MirSpan,
     ) -> Result<()> {
         let func = args.get(0).ok_or(MirEvalError::TypeError("fn trait with no arg"))?;
@@ -1958,7 +1960,38 @@ impl Evaluator<'_> {
                     span,
                 )?;
             }
-            x => not_supported!("Call FnTrait methods with type {x:?}"),
+            _ => {
+                // try to execute the manual impl of `FnTrait` for structs (nightly feature used in std)
+                let arg0 = func;
+                let args = &args[1..];
+                let arg1 = {
+                    let ty = TyKind::Tuple(
+                        args.len(),
+                        Substitution::from_iter(Interner, args.iter().map(|x| x.ty.clone())),
+                    )
+                    .intern(Interner);
+                    let layout = self.layout(&ty)?;
+                    let result = self.make_by_layout(
+                        layout.size.bytes_usize(),
+                        &layout,
+                        None,
+                        args.iter().map(|x| IntervalOrOwned::Borrowed(x.interval)),
+                    )?;
+                    // FIXME: there is some leak here
+                    let size = layout.size.bytes_usize();
+                    let addr = self.heap_allocate(size, layout.align.abi.bytes() as usize);
+                    self.write_memory(addr, &result)?;
+                    IntervalAndTy { interval: Interval { addr, size }, ty }
+                };
+                return self.exec_fn_with_args(
+                    def,
+                    &[arg0.clone(), arg1],
+                    generic_args,
+                    locals,
+                    destination,
+                    span,
+                );
+            }
         }
         Ok(())
     }
diff --git a/crates/hir-ty/src/mir/eval/shim.rs b/crates/hir-ty/src/mir/eval/shim.rs
index 3b9ef03c369..3ed26d0e9ba 100644
--- a/crates/hir-ty/src/mir/eval/shim.rs
+++ b/crates/hir-ty/src/mir/eval/shim.rs
@@ -681,7 +681,23 @@ impl Evaluator<'_> {
                     let addr = tuple.interval.addr.offset(offset);
                     args.push(IntervalAndTy::new(addr, field, self, locals)?);
                 }
-                self.exec_fn_trait(&args, destination, locals, span)
+                if let Some(target) = self.db.lang_item(self.crate_id, LangItem::FnOnce) {
+                    if let Some(def) = target
+                        .as_trait()
+                        .and_then(|x| self.db.trait_data(x).method_by_name(&name![call_once]))
+                    {
+                        return self.exec_fn_trait(
+                            def,
+                            &args,
+                            // FIXME: wrong for manual impls of `FnOnce`
+                            Substitution::empty(Interner),
+                            locals,
+                            destination,
+                            span,
+                        );
+                    }
+                }
+                not_supported!("FnOnce was not available for executing const_eval_select");
             }
             _ => not_supported!("unknown intrinsic {name}"),
         }