charon_driver/translate/
translate_closures.rs

1//! In rust, closures behave like ADTs that implement the FnOnce/FnMut/Fn traits automatically.
2//! Function item types also implement those traits automatically.
3//!
4//! Here we convert closures to a struct containing the closure's state (upvars), along with
5//! matching trait impls and fun decls (e.g. a Fn closure will have a trait impl for Fn, FnMut and
6//! FnOnce, along with 3 matching method implementations for call, call_mut and call_once).
7//! Function item types reuse the same generated trait impls and methods, with an empty state that
8//! forwards to the original function item.
9//!
10//! For example, given the following Rust code:
11//! ```ignore
12//! pub fn test_closure_capture<T: Clone>() {
13//!     let mut v = vec![];
14//!     let mut add = |x: &u32| v.push(*x);
15//!     add(&0);
16//!     add(&1);
17//! }
18//! ```
19//!
20//! We generate the equivalent desugared code:
21//! ```text
22//! struct {test_closure_capture::closure#0}<'a, T: Clone> (&'a mut Vec<u32>);
23//!
24//! // The 'a comes from captured variables, the 'b comes from the closure higher-kinded signature.
25//! impl<'a, 'b, T: Clone> FnMut<(&'b u32,)> for {test_closure_capture::closure#0}<'a, T> {
26//!     fn call_mut<'c>(&'c mut self, arg: (&'b u32,)) {
27//!         self.0.push(*arg.0);
28//!     }
29//! }
30//!
31//! impl<'a, 'b, T: Clone> FnOnce<(&'b u32,)> for {test_closure_capture::closure#0}<'a, T> {
32//!     type Output = ();
33//!     ...
34//! }
35//!
36//! pub fn test_closure_capture<T: Clone>() {
37//!     let mut v = vec![];
38//!     let mut add = {test_closure_capture::closure#0} (&mut v);
39//!     state.call_mut(&0);
40//!     state.call_mut(&1);
41//! }
42//! ```
43
44use crate::hax::{self, UnderOwnerState};
45use itertools::Itertools;
46use std::mem;
47
48use super::translate_crate::TransItemSourceKind;
49use super::translate_ctx::*;
50use charon_lib::ullbc_ast::*;
51
52pub fn translate_closure_kind(kind: &hax::ClosureKind) -> ClosureKind {
53    match kind {
54        hax::ClosureKind::Fn => ClosureKind::Fn,
55        hax::ClosureKind::FnMut => ClosureKind::FnMut,
56        hax::ClosureKind::FnOnce => ClosureKind::FnOnce,
57    }
58}
59
60/// If this trait proof is a built-in impl of a `Fn*` trait, return the `Self` type it is
61/// implemented for and the kind of the implemented trait.
62pub fn recognize_fn_trait_impl_proof<'a, 'tcx>(
63    trait_proof: &'a hax::TraitProof<'tcx>,
64) -> Option<(hax::Binder<&'a hax::Ty<'tcx>>, ClosureKind)> {
65    let hax::TraitProofKind::Builtin {
66        trait_data: hax::BuiltinTraitData::Other(lang_item),
67        ..
68    } = &trait_proof.kind
69    else {
70        return None;
71    };
72    let kind = match lang_item {
73        hax::SolverTraitLangItem::FnOnce => ClosureKind::FnOnce,
74        hax::SolverTraitLangItem::FnMut => ClosureKind::FnMut,
75        hax::SolverTraitLangItem::Fn => ClosureKind::Fn,
76        _ => return None,
77    };
78    let Some(hax::GenericArg::Type(self_ty)) =
79        trait_proof.pred.hax_skip_binder_ref().generic_args.first()
80    else {
81        unreachable!("no `Self` type arg on a `Fn*` trait ref")
82    };
83    Some((trait_proof.pred.rebind(self_ty), kind))
84}
85
86/// The built-in `Fn*` impl of the given kind that we generate for this closure or function item.
87pub fn callable_virtual_impl<'a, 'tcx>(
88    s: &impl hax::BaseState<'tcx>,
89    def: &'a hax::FullDef<'tcx>,
90    target_kind: ClosureKind,
91) -> &'a hax::VirtualTraitImpl<'tcx> {
92    CallableFnImpls::from_def(s, def, target_kind)
93        .and_then(|impls| impls.vimpl)
94        .expect("expected a callable with a Fn* impl")
95}
96
97#[derive(Copy, Clone)]
98enum CallableKind {
99    Closure(ClosureKind),
100    FnDef,
101    FnPtr,
102}
103
104#[derive(Copy, Clone)]
105struct Callable<'tcx, 'a> {
106    kind: CallableKind,
107    item: &'a hax::ItemRef<'tcx>,
108    sig: &'a hax::PolyFnSig<'tcx>,
109}
110
111impl<'tcx, 'a> Callable<'tcx, 'a> {
112    fn for_closure(
113        s: &impl hax::BaseState<'tcx>,
114        closure: &'a hax::ClosureTy<'tcx>,
115        sig: &'a hax::PolyFnSig<'tcx>,
116    ) -> Self {
117        Self {
118            item: &closure.item,
119            kind: CallableKind::Closure(translate_closure_kind(&closure.kind(s))),
120            sig,
121        }
122    }
123
124    fn for_closure_def(def: &'a hax::FullDef<'tcx>) -> Self {
125        let hax::FullDefKind::Closure(closure) = def.kind() else {
126            unreachable!()
127        };
128        Self {
129            item: def.this(),
130            kind: CallableKind::Closure(translate_closure_kind(closure.kind())),
131            sig: closure.sig(),
132        }
133    }
134
135    fn kind(self) -> CallableKind {
136        self.kind
137    }
138
139    fn item(self) -> &'a hax::ItemRef<'tcx> {
140        self.item
141    }
142
143    fn sig(self) -> &'a hax::PolyFnSig<'tcx> {
144        self.sig
145    }
146}
147
148#[derive(Copy, Clone)]
149struct CallableFnImpls<'a, 'tcx> {
150    callable: Callable<'tcx, 'a>,
151    /// The virtual impl of the requested `Fn*` trait, if the callable implements it.
152    vimpl: Option<&'a hax::VirtualTraitImpl<'tcx>>,
153}
154
155impl<'a, 'tcx> CallableFnImpls<'a, 'tcx> {
156    fn from_def(
157        s: &impl hax::BaseState<'tcx>,
158        def: &'a hax::FullDef<'tcx>,
159        target_kind: ClosureKind,
160    ) -> Option<Self> {
161        let vimpl_from_impls = |impls: &'a hax::FnTraitImpls<'tcx>| -> Option<_> {
162            match target_kind {
163                ClosureKind::FnOnce => Some(impls.fn_once_impl.as_ref()),
164                ClosureKind::FnMut => impls.fn_mut_impl.as_deref(),
165                ClosureKind::Fn => impls.fn_impl.as_deref(),
166            }
167        };
168        let for_fn_def = |sig: &'a hax::PolyFnSig<'tcx>, impls: &'a hax::FnTraitImpls<'tcx>| {
169            Some(Self {
170                callable: Callable {
171                    kind: CallableKind::FnDef,
172                    item: def.this(),
173                    sig,
174                },
175                vimpl: vimpl_from_impls(impls),
176            })
177        };
178        match def.kind() {
179            hax::FullDefKind::Closure(c) => Some(Self {
180                callable: Callable::for_closure_def(def),
181                vimpl: vimpl_from_impls(c.fn_trait_impls(s)),
182            }),
183            hax::FullDefKind::Fn(f) => for_fn_def(f.sig(), f.fn_trait_impls(s)?),
184            hax::FullDefKind::AssocFn(f) => for_fn_def(f.sig(), f.fn_trait_impls(s)?),
185            hax::FullDefKind::Ctor(f) => for_fn_def(f.sig(), f.fn_trait_impls(s)?),
186            hax::FullDefKind::FnPtr(f) => {
187                let impls = f.fn_trait_impls(s)?;
188                let sig = f.sig();
189                Some(Self {
190                    callable: Callable {
191                        kind: CallableKind::FnPtr,
192                        item: def.this(),
193                        sig,
194                    },
195                    vimpl: vimpl_from_impls(impls),
196                })
197            }
198            _ => None,
199        }
200    }
201}
202
203/// References to callable items are subtle because there are three sources of lifetimes on top of
204/// the normal generics: closure upvars, the higher-kindedness of the callable itself, and the
205/// late-bound generics of the `call`/`call_mut` methods. One must be careful to choose the right
206/// method from these.
207impl<'tcx> ItemTransCtx<'tcx, '_> {
208    /// Translate a reference to a callable item that takes late-bound lifetimes. The binder binds
209    /// the late-bound lifetimes of the callable itself, if it is higher-kinded.
210    fn translate_callable_bound_ref_with_late_bound(
211        &mut self,
212        span: Span,
213        callable: Callable<'tcx, '_>,
214        kind: TransItemSourceKind,
215    ) -> Result<RegionBinder<DeclRef<ItemId>>, Error> {
216        if !matches!(
217            kind,
218            TransItemSourceKind::TraitImpl(..) | TransItemSourceKind::ClosureAsFnCast
219        ) {
220            raise_error!(
221                self,
222                span,
223                "Called `translate_callable_bound_ref_with_late_bound` on a `{kind:?}`; \
224                use `translate_closure_ref_with_upvars` \
225                or `translate_callable_bound_ref_with_method_bound` instead"
226            )
227        }
228        let dref: DeclRef<ItemId> = self.translate_item(span, callable.item(), kind)?;
229        self.translate_region_binder(span, callable.sig(), |ctx, _| {
230            let mut dref = dref.move_under_binder();
231            // The regions for these item kinds have the fn late bound regions at the end.
232            for (a, b) in dref.generics.regions.iter_mut().rev().zip(
233                ctx.innermost_binder()
234                    .params
235                    .identity_args()
236                    .regions
237                    .into_iter()
238                    .rev(),
239            ) {
240                *a = b;
241            }
242            Ok(dref)
243        })
244    }
245
246    /// Translate a reference to a callable item that takes late-bound lifetimes and method
247    /// lifetimes. The binder binds the late-bound lifetimes of the `call`/`call_mut` method
248    /// (specified by `target_kind`).
249    fn translate_callable_bound_ref_with_method_bound(
250        &mut self,
251        span: Span,
252        item: &hax::ItemRef<'tcx>,
253        kind: TransItemSourceKind,
254        target_kind: ClosureKind,
255    ) -> Result<RegionBinder<DeclRef<ItemId>>, Error> {
256        if !matches!(kind, TransItemSourceKind::CallableMethod(..)) {
257            raise_error!(
258                self,
259                span,
260                "Called `translate_callable_bound_ref_with_method_bound` on a `{kind:?}`; \
261                use `translate_closure_ref_with_upvars` \
262                or `translate_callable_bound_ref_with_late_bound` instead"
263            )
264        }
265        let dref: DeclRef<ItemId> = self.translate_item(span, item, kind)?;
266        let mut dref = dref.move_under_binder();
267        let mut regions = IndexVec::new();
268        match target_kind {
269            ClosureKind::FnOnce => {}
270            ClosureKind::FnMut | ClosureKind::Fn => {
271                let rid =
272                    regions.push_with(|index| RegionParam::new(index, None, Variance::Covariant));
273                *dref.generics.regions.iter_mut().last().unwrap() =
274                    Region::Var(DeBruijnVar::new_at_zero(rid));
275            }
276        }
277        Ok(RegionBinder {
278            regions,
279            outlives: Vec::new(),
280            skip_binder: dref,
281        })
282    }
283
284    /// If this trait proof is the built-in impl of a `Fn*` trait for a closure, function item or
285    /// function pointer, return the callable item and the kind of the implemented trait. The
286    /// returned item has its regions erased.
287    pub fn recognize_callable_impl_proof(
288        &self,
289        trait_proof: &hax::TraitProof<'tcx>,
290    ) -> Option<(hax::ItemRef<'tcx>, ClosureKind)> {
291        let hax::TraitProofKind::Builtin {
292            trait_data: hax::BuiltinTraitData::Other(lang_item),
293            ..
294        } = &trait_proof.kind
295        else {
296            return None;
297        };
298        let target_kind = match lang_item {
299            hax::SolverTraitLangItem::FnOnce => ClosureKind::FnOnce,
300            hax::SolverTraitLangItem::FnMut => ClosureKind::FnMut,
301            hax::SolverTraitLangItem::Fn => ClosureKind::Fn,
302            _ => return None,
303        };
304        // We skip the binder and erase regions to avoid bound vars escaping.
305        let hax::GenericArg::Type(callable_ty) = trait_proof
306            .pred
307            .hax_skip_binder_ref()
308            .generic_args
309            .first()?
310        else {
311            return None;
312        };
313        let item = callable_ty.callable_item()?;
314        Some((item.erase(self.hax_state_with_id()), target_kind))
315    }
316
317    pub(crate) fn translate_callable_method_fn_ptr(
318        &mut self,
319        span: Span,
320        item: &hax::ItemRef<'tcx>,
321    ) -> Result<Option<RegionBinder<FnPtr>>, Error> {
322        if !self.monomorphize() {
323            return Ok(None);
324        }
325        let Some(in_trait) = &item.in_trait else {
326            return Ok(None);
327        };
328        let Some((callable, target_kind)) = self.recognize_callable_impl_proof(in_trait) else {
329            return Ok(None);
330        };
331        let kind = TransItemSourceKind::CallableMethod(target_kind);
332        let bound_ref = self.translate_callable_bound_ref_with_method_bound(
333            span,
334            &callable,
335            kind,
336            target_kind,
337        )?;
338        Ok(Some(bound_ref.map(|dref| {
339            let fn_ref: FunDeclRef = dref.try_into().unwrap();
340            FnPtr::new(FnPtrKind::Fun(fn_ref.id), fn_ref.generics)
341        })))
342    }
343}
344
345impl<'tcx> ItemTransCtx<'tcx, '_> {
346    /// Translate a reference to the closure ADT.
347    pub fn translate_closure_type_ref(
348        &mut self,
349        span: Span,
350        closure: &hax::ItemRef<'tcx>,
351    ) -> Result<TypeDeclRef, Error> {
352        self.translate_type_decl_ref(span, closure)
353    }
354
355    /// For stateless closures, translate a function reference to the top-level function that
356    /// executes the closure code without taking the state as parameter.If you want to instantiate
357    /// the binder, use the lifetimes from `self.closure_late_regions`.
358    pub fn translate_stateless_closure_as_fn_ref(
359        &mut self,
360        span: Span,
361        closure: &hax::ClosureTy<'tcx>,
362    ) -> Result<RegionBinder<FunDeclRef>, Error> {
363        let kind = TransItemSourceKind::ClosureAsFnCast;
364        let sig = closure.fn_sig(self.hax_state());
365        let callable = Callable::for_closure(self.hax_state(), closure, &sig);
366        let bound_dref = self.translate_callable_bound_ref_with_late_bound(span, callable, kind)?;
367        Ok(bound_dref.map(|dref| dref.try_into().unwrap()))
368    }
369
370    /// Translate a reference to the chosen closure impl. The resulting value needs lifetime
371    /// arguments for late-bound lifetimes. If you want to instantiate the binder, use the
372    /// lifetimes from `self.closure_late_regions`.
373    pub fn translate_closure_bound_impl_ref(
374        &mut self,
375        span: Span,
376        def: &hax::FullDef<'tcx>,
377        target_kind: ClosureKind,
378    ) -> Result<RegionBinder<TraitImplRef>, Error> {
379        let kind = TransItemSourceKind::TraitImpl(TransImplSource::Callable(target_kind));
380        let bound_dref = self.translate_callable_bound_ref_with_late_bound(
381            span,
382            Callable::for_closure_def(def),
383            kind,
384        )?;
385        Ok(bound_dref.map(|dref| dref.try_into().unwrap()))
386    }
387
388    /// Translate a reference to the chosen callable impl.
389    pub fn translate_callable_impl_ref(
390        &mut self,
391        span: Span,
392        item: &hax::ItemRef<'tcx>,
393        target_kind: ClosureKind,
394    ) -> Result<TraitImplRef, Error> {
395        self.translate_item(
396            span,
397            item,
398            TransItemSourceKind::TraitImpl(TransImplSource::Callable(target_kind)),
399        )
400    }
401
402    pub fn translate_closure_info(
403        &mut self,
404        span: Span,
405        def: &hax::FullDef<'tcx>,
406    ) -> Result<ClosureInfo, Error> {
407        use ClosureKind::*;
408        let hax::FullDefKind::Closure(closure) = def.kind() else {
409            unreachable!()
410        };
411        let kind = translate_closure_kind(closure.kind());
412
413        let closure_body =
414            self.translate_item(span, def.this(), TransItemSourceKind::CallableMethod(kind))?;
415
416        let fn_once_impl = self.translate_closure_bound_impl_ref(span, def, FnOnce)?;
417        let fn_mut_impl = if matches!(kind, FnMut | Fn) {
418            Some(self.translate_closure_bound_impl_ref(span, def, FnMut)?)
419        } else {
420            None
421        };
422        let fn_impl = if matches!(kind, Fn) {
423            Some(self.translate_closure_bound_impl_ref(span, def, Fn)?)
424        } else {
425            None
426        };
427        let signature = self.translate_poly_fun_sig(span, closure.sig())?;
428        Ok(ClosureInfo {
429            kind,
430            closure_body,
431            fn_once_impl,
432            fn_mut_impl,
433            fn_impl,
434            signature,
435        })
436    }
437
438    fn get_callable_state_ty(
439        &mut self,
440        span: Span,
441        callable: Callable<'tcx, '_>,
442    ) -> Result<Ty, Error> {
443        Ok(match callable.kind() {
444            CallableKind::Closure(..) => {
445                let tref = self.translate_closure_type_ref(span, callable.item())?;
446                TyKind::Adt(tref).into_ty()
447            }
448            CallableKind::FnDef => {
449                let fn_ref =
450                    self.translate_bound_fn_ptr(span, callable.item(), TransItemSourceKind::Fun)?;
451                TyKind::FnDef(fn_ref).into_ty()
452            }
453            CallableKind::FnPtr => {
454                TyKind::FnPtr(self.translate_poly_fun_sig(span, callable.sig())?).into_ty()
455            }
456        })
457    }
458
459    pub fn translate_closure_adt(
460        &mut self,
461        span: Span,
462        closure: &hax::Closure<'tcx>,
463    ) -> Result<TypeDeclKind, Error> {
464        let fields = closure
465            .upvar_tys()
466            .iter()
467            .enumerate()
468            .map(|(field_id, ty)| -> Result<_, Error> {
469                Ok(Field {
470                    span,
471                    attr_info: AttrInfo::dummy_private(),
472                    name: format!("_{field_id}"),
473                    is_positional: true,
474                    ty: self.translate_ty(span, ty)?,
475                })
476            })
477            .try_collect()?;
478        Ok(TypeDeclKind::Struct(fields))
479    }
480
481    fn translate_callable_method_body(
482        &mut self,
483        span: Span,
484        def: &hax::FullDef<'tcx>,
485        target_kind: ClosureKind,
486        callable: Callable<'tcx, '_>,
487        signature: &FunSig,
488    ) -> Result<Body, Error> {
489        match callable.kind() {
490            CallableKind::Closure(source_kind) => self.translate_closure_method_body(
491                span,
492                def,
493                source_kind,
494                target_kind,
495                callable.item(),
496                signature,
497            ),
498            CallableKind::FnDef => {
499                self.translate_fn_def_method_body(span, callable.item(), signature)
500            }
501            CallableKind::FnPtr => {
502                self.translate_fn_ptr_method_body(span, callable.item(), signature)
503            }
504        }
505    }
506
507    fn translate_closure_method_body(
508        &mut self,
509        span: Span,
510        def: &hax::FullDef<'tcx>,
511        source_kind: ClosureKind,
512        target_kind: ClosureKind,
513        closure: &hax::ItemRef<'tcx>,
514        signature: &FunSig,
515    ) -> Result<Body, Error> {
516        use ClosureKind::*;
517        Ok(match (target_kind, source_kind) {
518            (Fn, Fn) | (FnMut, FnMut) | (FnOnce, FnOnce) => {
519                // Translate the function's body normally
520                let mut body = self.translate_def_body(span, def);
521                // The body is translated as if the locals are: ret value, state, arg-1,
522                // ..., arg-N, rest...
523                // However, there is only one argument with the tupled closure arguments;
524                // we must thus shift all locals with index >=2 by 1, and add a new local
525                // for the tupled arg, giving us: ret value, state, args, arg-1, ...,
526                // arg-N, rest...
527                // We then add N statements of the form `locals[N+3] := move locals[2].N`,
528                // to destructure the arguments.
529                let Body::Unstructured(inner_body) = &mut body else {
530                    return Ok(body);
531                };
532
533                inner_body.dyn_visit_in_body_mut(|local: &mut LocalId| {
534                    if local.index() >= 2 {
535                        *local += 1;
536                    }
537                });
538
539                let GExprBody {
540                    locals,
541                    body: blocks,
542                    ..
543                } = inner_body;
544
545                // The (Arg1, Arg2, ..) type.
546                let tupled_ty = &signature.inputs[1];
547
548                // Remember how many arguments there are
549                let closure_arg_count = locals.arg_count - 1;
550                let mut old_locals = mem::take(&mut locals.locals).into_iter();
551                locals.arg_count = 2;
552                locals.locals.push(old_locals.next().unwrap()); // ret
553                locals.locals.push(old_locals.next().unwrap()); // state
554                let tupled_arg = locals.new_var(Some("tupled_args".to_string()), tupled_ty.clone());
555                locals.locals.extend(old_locals);
556
557                let untupled_args = locals
558                    .locals
559                    .iter()
560                    .skip(3)
561                    .take(closure_arg_count)
562                    .map(|l| &l.ty)
563                    .cloned();
564                let new_stts = untupled_args.enumerate().map(|(i, ty)| {
565                    let nth_field = tupled_arg
566                        .clone()
567                        .project(ProjectionElem::Field(None, FieldId::new(i)), ty);
568                    let local_id = LocalId::new(i + 3);
569                    Statement::new(
570                        span,
571                        StatementKind::Assign(
572                            locals.place_for_var(local_id),
573                            Rvalue::Use(Operand::Move(nth_field), WithRetag::No),
574                        ),
575                    )
576                });
577                blocks[BlockId::ZERO].statements.splice(0..0, new_stts);
578
579                body
580            }
581            // Target translation:
582            //
583            // fn call_once(state: Self, args: Args) -> Output {
584            //   let temp_ref = &[mut] state;
585            //   let ret = self.call[_mut](temp, args);
586            //   drop state;
587            //   return ret;
588            // }
589            //
590            (FnOnce, Fn | FnMut) => {
591                // Hax (via rustc) gives us the MIR to do this.
592                let Some(body) = def.this.closure_once_shim(self.hax_state()) else {
593                    panic!("missing shim for closure")
594                };
595                let mut body = self.translate_body(span, body, &def.source_text);
596                if let Some(body) = body.as_unstructured_mut() {
597                    // The shim generated by rustc doesn't declare its receiver mutable.
598                    body.locals.locals[LocalId::new(1)].mutable = true;
599                }
600                body
601            }
602            // Target translation:
603            //
604            // fn call_mut(state: &mut Self, args: Args) -> Output {
605            //   let reborrow = &*state;
606            //   self.call(reborrow, args)
607            // }
608            (FnMut, Fn) => {
609                let fun_id: FunDeclId = self.register_item(
610                    span,
611                    def.this(),
612                    TransItemSourceKind::CallableMethod(source_kind),
613                );
614                let impl_ref = self.translate_callable_impl_ref(span, closure, source_kind)?;
615                // TODO: make a trait call to avoid needing to concatenate things ourselves.
616                // TODO: can we ask hax for the trait ref?
617                let fn_op = FnOperand::Regular(FnPtr::new(
618                    fun_id.into(),
619                    impl_ref.generics.concat(&GenericArgs {
620                        regions: vec![self.translate_erased_region()].into(),
621                        ..GenericArgs::empty()
622                    }),
623                ));
624
625                let mut builder = BodyBuilder::new(span, 2);
626
627                let output = builder.new_var(None, signature.output.clone());
628                let state = builder.new_var(Some("state".to_string()), signature.inputs[0].clone());
629                let args = builder.new_var(Some("args".to_string()), signature.inputs[1].clone());
630                let deref_state = state.deref();
631                let reborrow_ty = TyKind::Ref(
632                    self.translate_erased_region(),
633                    deref_state.ty.clone(),
634                    RefKind::Shared,
635                )
636                .into_ty();
637                let reborrow = builder.new_var(None, reborrow_ty);
638
639                builder.push_statement(StatementKind::Assign(
640                    reborrow.clone(),
641                    Rvalue::Ref {
642                        place: deref_state,
643                        kind: BorrowKind::Shared,
644                        // The state must be Sized, hence `()` as ptr-metadata
645                        ptr_metadata: Operand::mk_const_unit(),
646                    },
647                ));
648
649                builder.call(
650                    output,
651                    Call {
652                        func: fn_op,
653                        args: vec![Operand::Move(reborrow), Operand::Move(args)],
654                        safety: CallSafety::Inherit,
655                    },
656                );
657
658                Body::Unstructured(builder.build())
659            }
660            (Fn, FnOnce) | (Fn, FnMut) | (FnMut, FnOnce) => {
661                panic!(
662                    "Can't make a closure body for a more restrictive kind \
663                    than the closure kind"
664                )
665            }
666        })
667    }
668
669    fn translate_fn_def_method_body(
670        &mut self,
671        span: Span,
672        item: &hax::ItemRef<'tcx>,
673        signature: &FunSig,
674    ) -> Result<Body, Error> {
675        let late_bound_regions: RegionArgs = self
676            .innermost_binder()
677            .bound_region_vars
678            .iter()
679            .map(|rid| Region::Var(DeBruijnVar::new_at_zero(*rid)))
680            .collect();
681        let fn_ptr = self
682            .translate_bound_fn_ptr(span, item, TransItemSourceKind::Fun)?
683            .apply(late_bound_regions);
684        let fn_op = FnOperand::Regular(fn_ptr);
685
686        let mut builder = BodyBuilder::new(span, 2);
687
688        let output = builder.new_var(None, signature.output.clone());
689        let _state = builder.new_var(Some("state".to_string()), signature.inputs[0].clone());
690        let tupled_args_ty = &signature.inputs[1];
691        let tupled_args = builder.new_var(Some("args".to_string()), tupled_args_ty.clone());
692
693        // We need the type declaration to have been translated to get the fields (since in monomorphic)
694        // mode they aren't in the generics. So ensure it's been translated!
695        let tuple_id = tupled_args_ty.as_adt().unwrap().id;
696        let _ = self.get_or_translate(ItemId::Type(tuple_id))?;
697        let arg_tys = tupled_args_ty.as_tuple_fields(&self.t_ctx.translated);
698
699        let args = arg_tys
700            .into_iter()
701            .enumerate()
702            .map(|(i, ty)| {
703                let nth_field = tupled_args
704                    .clone()
705                    .project(ProjectionElem::Field(None, FieldId::new(i)), ty);
706                Operand::Move(nth_field)
707            })
708            .collect();
709
710        builder.call(
711            output,
712            Call {
713                func: fn_op,
714                args,
715                safety: CallSafety::Inherit,
716            },
717        );
718
719        Ok(Body::Unstructured(builder.build()))
720    }
721
722    fn translate_fn_ptr_method_body(
723        &mut self,
724        span: Span,
725        item: &hax::ItemRef<'tcx>,
726        signature: &FunSig,
727    ) -> Result<Body, Error> {
728        let ty = item
729            .def_id
730            .type_of(&self.hax_state)
731            .instantiate(self.tcx, item.rustc_args(&self.hax_state));
732        let ty = hax::normalize(self.tcx, self.hax_state.typing_env(), ty);
733        let copy_proof =
734            hax::solve_copy(&self.hax_state, ty).expect("function pointers are always Copy");
735        let copy_proof = self
736            .translate_trait_proof(span, &copy_proof)?
737            .no_bound_vars();
738        let mut builder = BodyBuilder::new(span, 2);
739        let output = builder.new_var(None, signature.output.clone());
740        let fn_ptr = builder.new_var(Some("fn_ptr".to_string()), signature.inputs[0].clone());
741        let tupled_args_ty = &signature.inputs[1];
742        let tupled_args = builder.new_var(Some("args".to_string()), tupled_args_ty.clone());
743
744        // `FnOnce::call_once` receives the function pointer directly, whereas `Fn::call` and
745        // `FnMut::call_mut` receive it through a shared or mutable reference respectively.
746        let fn_ptr = if matches!(fn_ptr.ty().kind(), TyKind::FnPtr(..)) {
747            fn_ptr
748        } else {
749            fn_ptr.deref()
750        };
751        let fn_ptr = Operand::Copy(fn_ptr, copy_proof);
752        let TyKind::FnPtr(fn_sig) = fn_ptr.ty().kind() else {
753            unreachable!()
754        };
755        let late_bound_regions: RegionArgs = self
756            .innermost_binder()
757            .bound_region_vars
758            .iter()
759            .map(|rid| Region::Var(DeBruijnVar::new_at_zero(*rid)))
760            .collect();
761        let arg_tys = fn_sig.clone().apply(late_bound_regions.clone()).inputs;
762        let args = arg_tys
763            .into_iter()
764            .enumerate()
765            .map(|(i, ty)| {
766                let field = tupled_args
767                    .clone()
768                    .project(ProjectionElem::Field(None, FieldId::new(i)), ty);
769                Operand::Move(field)
770            })
771            .collect();
772        builder.call(
773            output,
774            Call {
775                func: FnOperand::Dynamic(fn_ptr, late_bound_regions),
776                args,
777                safety: CallSafety::Inherit,
778            },
779        );
780        Ok(Body::Unstructured(builder.build()))
781    }
782
783    /// Given an item that is a closure, generate the `call_once`/`call_mut`/`call` method
784    /// (depending on `target_kind`).
785    #[tracing::instrument(skip(self, item_meta))]
786    pub fn translate_closure_method(
787        mut self,
788        def_id: FunDeclId,
789        item_meta: ItemMeta,
790        def: &hax::FullDef<'tcx>,
791        target_kind: ClosureKind,
792    ) -> Result<FunDecl, Error> {
793        let span = item_meta.span;
794        let callable_impls = CallableFnImpls::from_def(self.hax_state(), def, target_kind).unwrap();
795        let callable = callable_impls.callable;
796
797        // Hax gives us trait-related information for the impl we're building.
798        let vimpl = callable_impls.vimpl.unwrap();
799        let implemented_trait = self.translate_trait_predicate(span, &vimpl.trait_pred)?;
800        let method = &vimpl.methods[0];
801        let method_id = self.translate_trait_method_id(implemented_trait.id, &method.def_id)?;
802
803        let impl_ref = self.translate_callable_impl_ref(span, callable.item(), target_kind)?;
804        let src = FunSource::TraitImpl {
805            impl_ref,
806            trait_ref: implemented_trait.clone(),
807            item_id: method_id,
808            reuses_default: false,
809        };
810
811        // Translate the function signature
812        let bound_sig = self.translate_poly_fun_sig(span, &method.sig)?;
813        // We give it the lifetime parameter we had prepared for that purpose.
814        let signature = bound_sig.apply(
815            self.the_only_binder()
816                .closure_call_method_region
817                .iter()
818                .map(|r| Region::Var(DeBruijnVar::new_at_zero(*r)))
819                .collect(),
820        );
821
822        let body = if item_meta.opacity.with_private_contents().is_opaque() {
823            Body::Opaque
824        } else {
825            self.translate_callable_method_body(span, def, target_kind, callable, &signature)?
826        };
827
828        Ok(FunDecl {
829            def_id,
830            item_meta,
831            generics: self.into_generics(),
832            signature: Box::new(signature),
833            src,
834            body,
835        })
836    }
837
838    #[tracing::instrument(skip(self, item_meta))]
839    pub fn translate_closure_trait_impl(
840        mut self,
841        def_id: TraitImplId,
842        item_meta: ItemMeta,
843        def: &hax::FullDef<'tcx>,
844        target_kind: ClosureKind,
845    ) -> Result<TraitImpl, Error> {
846        let span = item_meta.span;
847        let callable_impls = CallableFnImpls::from_def(self.hax_state(), def, target_kind).unwrap();
848        let callable = callable_impls.callable;
849
850        // Hax gives us trait-related information for the impl we're building.
851        let vimpl = callable_impls.vimpl.unwrap();
852        let mut timpl = self.translate_virtual_trait_impl(
853            def_id,
854            item_meta,
855            callable.item(),
856            TransImplSource::Callable(target_kind),
857            vimpl,
858        )?;
859
860        if self.monomorphize() {
861            return Ok(timpl);
862        }
863
864        // Construct the `call_*` method reference.
865        let trait_decl_id = timpl.impl_trait.id;
866        let trait_method_id =
867            self.translate_trait_method_id(trait_decl_id, &vimpl.methods[0].def_id)?;
868        let call_fn_binder = {
869            let kind = TransItemSourceKind::CallableMethod(target_kind);
870            let bound_method_ref: RegionBinder<DeclRef<ItemId>> = self
871                .translate_callable_bound_ref_with_method_bound(
872                    span,
873                    callable.item(),
874                    kind,
875                    target_kind,
876                )?;
877            let params = GenericParams {
878                regions: bound_method_ref.regions,
879                ..GenericParams::empty()
880            };
881            let fn_decl_ref: FunDeclRef = bound_method_ref.skip_binder.try_into().unwrap();
882            Binder::new(
883                BinderKind::TraitMethod(trait_decl_id, trait_method_id),
884                params,
885                fn_decl_ref,
886            )
887        };
888        timpl
889            .methods
890            .set_slot_extend(trait_method_id, call_fn_binder);
891
892        Ok(timpl)
893    }
894
895    /// Given an item that is a non-capturing closure, generate the equivalent function,
896    /// by removing the state from the parameters and untupling the arguments.
897    #[tracing::instrument(skip(self, item_meta))]
898    pub fn translate_stateless_closure_as_fn(
899        mut self,
900        def_id: FunDeclId,
901        item_meta: ItemMeta,
902        def: &hax::FullDef<'tcx>,
903    ) -> Result<FunDecl, Error> {
904        let span = item_meta.span;
905        let hax::FullDefKind::Closure(closure) = &def.kind else {
906            unreachable!()
907        };
908
909        trace!("About to translate closure as fn:\n{:?}", def.def_id());
910
911        assert!(
912            closure.upvar_tys().is_empty(),
913            "Only stateless closures can be translated as functions"
914        );
915        // Translate the function signature
916        let signature = self.translate_fun_sig(span, closure.sig().hax_skip_binder_ref())?;
917        let state_ty = self.get_callable_state_ty(span, Callable::for_closure_def(def))?;
918
919        let body = if item_meta.opacity.with_private_contents().is_opaque() {
920            Body::Opaque
921        } else {
922            // Target translation:
923            //
924            // fn call_fn(arg0: Args[0], ..., argN: Args[N]) -> Output {
925            //   let closure: Closure = {};
926            //   let args = (arg0, ..., argN);
927            //   closure.call(args)
928            // }
929            let fun_id: FunDeclId = self.register_item(
930                span,
931                def.this(),
932                TransItemSourceKind::CallableMethod(ClosureKind::FnOnce),
933            );
934            let impl_ref =
935                self.translate_callable_impl_ref(span, def.this(), ClosureKind::FnOnce)?;
936            let fn_op = FnOperand::Regular(FnPtr::new(fun_id.into(), impl_ref.generics.clone()));
937
938            let mut builder = BodyBuilder::new(span, signature.inputs.len());
939
940            let output = builder.new_var(None, signature.output.clone());
941            let args: Vec<Place> = signature
942                .inputs
943                .iter()
944                .enumerate()
945                .map(|(i, ty)| builder.new_var(Some(format!("arg{}", i + 1)), ty.clone()))
946                .collect();
947            let args_tupled_ty =
948                self.translate_ty(span, &closure.sig().hax_skip_binder_ref().tupled_inputs)?;
949            let args_tupled = builder.new_var(Some("args".to_string()), args_tupled_ty.clone());
950            let state = builder.new_var(Some("state".to_string()), state_ty.clone());
951
952            builder.push_statement(StatementKind::Assign(
953                args_tupled.clone(),
954                Rvalue::Aggregate(
955                    AggregateKind::Adt(args_tupled_ty.as_adt().unwrap().clone(), None, None),
956                    args.into_iter().map(Operand::Move).collect(),
957                ),
958            ));
959
960            let state_ty_adt = state_ty.as_adt().unwrap();
961            builder.push_statement(StatementKind::Assign(
962                state.clone(),
963                Rvalue::Aggregate(AggregateKind::Adt(state_ty_adt.clone(), None, None), vec![]),
964            ));
965
966            builder.call(
967                output,
968                Call {
969                    func: fn_op,
970                    args: vec![Operand::Move(state), Operand::Move(args_tupled)],
971                    safety: CallSafety::Inherit,
972                },
973            );
974
975            Body::Unstructured(builder.build())
976        };
977
978        Ok(FunDecl {
979            def_id,
980            item_meta,
981            generics: self.into_generics(),
982            signature: Box::new(signature),
983            src: FunSource::Normal,
984            body,
985        })
986    }
987}
charon_driver/translate/translate_closures.rs

charon_driver/translate/
translate_closures.rs
